test: wait on the writer instead of guessing at it in the BigQuery plugin tests

27 tests logged an event and then slept for a fixed 50ms before asserting
on what the batch writer had written. The sleep was standing in for a
barrier, so each of these assertions held only while the runner stayed
faster than the guess, and the suite runs on shared runners under
xdist where it sometimes is not.

The plugin already exposes the barrier these tests want: flush() joins
the write queue, and the queue is only marked done after the write
attempt completes. Several tests in this file already used it. Use it
everywhere the sleep was standing in for it.

The sleeps that remain are doing something else: forcing an interleave
between concurrent tasks, simulating setup latency, or hanging a writer
on purpose to exercise a timeout. Those are left alone.

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 953909736

Author: GWeale

tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py

@@ -2226,7 +2226,7 @@ async def test_on_agent_error_callback_logs_correctly(
         callback_context=callback_context,
         error=error,
     )
-    await asyncio.sleep(0.05)
+    await bq_plugin_inst.flush()
     rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
     log_entry = next(r for r in rows if r["event_type"] == "AGENT_ERROR")
     assert log_entry["error_message"] == "Agent crashed"
@@ -2267,7 +2267,7 @@ async def test_on_agent_error_does_not_pop_foreign_invocation_span(
         callback_context=callback_context,
         error=error,
     )
-    await asyncio.sleep(0.05)
+    await bq_plugin_inst.flush()
 
     # The invocation root was NOT consumed by the agent-error pop.
     assert trace_manager.get_current_span_id() == inv_span_id
@@ -2297,7 +2297,7 @@ async def test_on_run_error_callback_logs_correctly(
         invocation_context=invocation_context,
         error=error,
     )
-    await asyncio.sleep(0.05)
+    await bq_plugin_inst.flush()
     rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
     log_entry = next(r for r in rows if r["event_type"] == "INVOCATION_ERROR")
     assert log_entry["error_message"] == "Invocation failed"
@@ -2376,7 +2376,7 @@ async def test_traceback_not_truncated_with_negative_max_len(
           ),
           error=error,
       )
-      await asyncio.sleep(0.05)
+      await plugin.flush()
       rows = await _get_captured_rows_async(
           mock_write_client, dummy_arrow_schema
       )
@@ -5698,7 +5698,7 @@ async def test_hitl_confirmation_emits_additional_event(
     await bq_plugin_inst.on_event_callback(
         invocation_context=invocation_context, event=event
     )
-    await asyncio.sleep(0.05)
+    await bq_plugin_inst.flush()
     rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
     event_types = [r["event_type"] for r in rows]
     assert "HITL_CONFIRMATION_REQUEST" in event_types
@@ -5715,7 +5715,7 @@ async def test_hitl_credential_emits_additional_event(
     await bq_plugin_inst.on_event_callback(
         invocation_context=invocation_context, event=event
     )
-    await asyncio.sleep(0.05)
+    await bq_plugin_inst.flush()
     rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
     event_types = [r["event_type"] for r in rows]
     assert "HITL_CREDENTIAL_REQUEST" in event_types
@@ -5732,7 +5732,7 @@ async def test_hitl_completion_emits_additional_event(
     await bq_plugin_inst.on_event_callback(
         invocation_context=invocation_context, event=event
     )
-    await asyncio.sleep(0.05)
+    await bq_plugin_inst.flush()
     rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
     event_types = [r["event_type"] for r in rows]
     assert "HITL_CONFIRMATION_REQUEST_COMPLETED" in event_types
@@ -5749,7 +5749,7 @@ async def test_regular_tool_no_hitl_event(
     await bq_plugin_inst.on_event_callback(
         invocation_context=invocation_context, event=event
     )
-    await asyncio.sleep(0.05)
+    await bq_plugin_inst.flush()
     # No HITL events should be emitted for non-HITL function calls.
     # on_event_callback only logs STATE_DELTA and HITL events; a regular
     # function call produces neither.
@@ -7903,7 +7903,7 @@ async def test_cache_metadata_logged_when_present(
         callback_context=callback_context,
         llm_response=llm_response,
     )
-    await asyncio.sleep(0.05)
+    await bq_plugin_inst.flush()
     rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
     log_entry = next(r for r in rows if r["event_type"] == "LLM_RESPONSE")
 
@@ -7938,7 +7938,7 @@ def __init__(self):
         callback_context=callback_context,
         llm_response=mock_response,
     )
-    await asyncio.sleep(0.05)
+    await bq_plugin_inst.flush()
     rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
     log_entry = next(r for r in rows if r["event_type"] == "LLM_RESPONSE")
 
@@ -7986,7 +7986,7 @@ async def test_a2a_interaction_logged_for_response_metadata(
     )
     assert result is None
 
-    await asyncio.sleep(0.05)
+    await bq_plugin_inst.flush()
     rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
     event_types = [r["event_type"] for r in rows]
     assert "A2A_INTERACTION" in event_types
@@ -8025,7 +8025,7 @@ async def test_a2a_interaction_logged_for_request_metadata(
     )
     assert result is None
 
-    await asyncio.sleep(0.05)
+    await bq_plugin_inst.flush()
     rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
     event_types = [r["event_type"] for r in rows]
     assert "A2A_INTERACTION" in event_types
@@ -8058,7 +8058,7 @@ async def test_no_a2a_interaction_for_irrelevant_metadata(
     )
     assert result is None
 
-    await asyncio.sleep(0.05)
+    await bq_plugin_inst.flush()
     # No events logged — a2a:task_id alone is not a meaningful
     # interaction payload.
     assert mock_write_client.append_rows.call_count == 0
@@ -8078,7 +8078,7 @@ async def test_no_a2a_interaction_for_no_metadata(
     )
     assert result is None
 
-    await asyncio.sleep(0.05)
+    await bq_plugin_inst.flush()
     assert mock_write_client.append_rows.call_count == 0
 
 
@@ -8392,7 +8392,7 @@ async def test_logs_final_text_response(
     await bq_plugin_inst.on_event_callback(
         invocation_context=invocation_context, event=event
     )
-    await asyncio.sleep(0.05)
+    await bq_plugin_inst.flush()
     rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
     agent_resp_rows = [r for r in rows if r["event_type"] == "AGENT_RESPONSE"]
     assert len(agent_resp_rows) == 1
@@ -8420,7 +8420,7 @@ async def test_skips_function_call_events(
     await bq_plugin_inst.on_event_callback(
         invocation_context=invocation_context, event=event
     )
-    await asyncio.sleep(0.05)
+    await bq_plugin_inst.flush()
     assert mock_write_client.append_rows.call_count == 0
 
   @pytest.mark.asyncio
@@ -8441,7 +8441,7 @@ async def test_skips_function_response_events(
     await bq_plugin_inst.on_event_callback(
         invocation_context=invocation_context, event=event
     )
-    await asyncio.sleep(0.05)
+    await bq_plugin_inst.flush()
     assert mock_write_client.append_rows.call_count == 0
 
   @pytest.mark.asyncio
@@ -8461,7 +8461,7 @@ async def test_skips_partial_events(
     await bq_plugin_inst.on_event_callback(
         invocation_context=invocation_context, event=event
     )
-    await asyncio.sleep(0.05)
+    await bq_plugin_inst.flush()
     assert mock_write_client.append_rows.call_count == 0
 
   @pytest.mark.asyncio
@@ -8488,7 +8488,7 @@ async def test_skips_long_running_tool_events(
     await bq_plugin_inst.on_event_callback(
         invocation_context=invocation_context, event=event
     )
-    await asyncio.sleep(0.05)
+    await bq_plugin_inst.flush()
     rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
     types_emitted = [r["event_type"] for r in rows]
     assert "AGENT_RESPONSE" not in types_emitted
@@ -8513,7 +8513,7 @@ async def test_skips_thought_only_events(
     await bq_plugin_inst.on_event_callback(
         invocation_context=invocation_context, event=event
     )
-    await asyncio.sleep(0.05)
+    await bq_plugin_inst.flush()
     assert mock_write_client.append_rows.call_count == 0
 
   @pytest.mark.asyncio
@@ -8539,7 +8539,7 @@ async def test_mixed_thought_and_visible_logs_only_visible(
     await bq_plugin_inst.on_event_callback(
         invocation_context=invocation_context, event=event
     )
-    await asyncio.sleep(0.05)
+    await bq_plugin_inst.flush()
     rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
     agent_resp_rows = [r for r in rows if r["event_type"] == "AGENT_RESPONSE"]
     assert len(agent_resp_rows) == 1
@@ -8563,7 +8563,7 @@ async def test_skips_empty_part_events(
     await bq_plugin_inst.on_event_callback(
         invocation_context=invocation_context, event=event
     )
-    await asyncio.sleep(0.05)
+    await bq_plugin_inst.flush()
     assert mock_write_client.append_rows.call_count == 0
 
   @pytest.mark.asyncio
@@ -8582,7 +8582,7 @@ async def test_skips_empty_text_events(
     await bq_plugin_inst.on_event_callback(
         invocation_context=invocation_context, event=event
     )
-    await asyncio.sleep(0.05)
+    await bq_plugin_inst.flush()
     assert mock_write_client.append_rows.call_count == 0
 
   @pytest.mark.asyncio
@@ -8609,7 +8609,7 @@ async def test_skips_executable_code_only_events(
     await bq_plugin_inst.on_event_callback(
         invocation_context=invocation_context, event=event
     )
-    await asyncio.sleep(0.05)
+    await bq_plugin_inst.flush()
     assert mock_write_client.append_rows.call_count == 0
 
 
@@ -9913,7 +9913,7 @@ async def test_final_attributes_pass_redacts_direct_producers(
               },
           ),
       )
-      await asyncio.sleep(0.01)
+      await plugin.flush()
       log_entry = await _get_captured_event_dict_async(
           mock_write_client, dummy_arrow_schema
       )
@@ -10701,7 +10701,7 @@ async def test_namedtuple_attribute_does_not_drop_row(
               extra_attributes={"point": Point(1, 2)},
           ),
       )
-      await asyncio.sleep(0.01)
+      await plugin.flush()
       log_entry = await _get_captured_event_dict_async(
           mock_write_client, dummy_arrow_schema
       )
@@ -10790,7 +10790,7 @@ async def test_depth_capped_payload_flags_row_truncated(
               extra_attributes={"deep": deep},
           ),
       )
-      await asyncio.sleep(0.01)
+      await plugin.flush()
       log_entry = await _get_captured_event_dict_async(
           mock_write_client, dummy_arrow_schema
       )