Merge branch 'main' into feat/app-name-override

Author: amitjborate

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
       )
@@ -10460,20 +10460,21 @@ async def fake_lazy_setup(**kwargs):
 
     def run_in_fresh_loop():
       try:
-        with mock.patch.object(
-            plugin, "_lazy_setup", side_effect=fake_lazy_setup
-        ):
-          asyncio.run(plugin._ensure_started())
+        asyncio.run(plugin._ensure_started())
       except BaseException as e:  # noqa: BLE001 - collecting for assertion
         errors.append(e)
 
-    threads = [
-        platform_thread.create_thread(run_in_fresh_loop) for _ in range(2)
-    ]
-    for t in threads:
-      t.start()
-    for t in threads:
-      t.join(timeout=10)
+    # Patch from this thread only. patch.object swaps a single shared
+    # attribute and is not itself thread safe, so entering it from both
+    # threads raced on _lazy_setup instead of on the code under test.
+    with mock.patch.object(plugin, "_lazy_setup", side_effect=fake_lazy_setup):
+      threads = [
+          platform_thread.create_thread(run_in_fresh_loop) for _ in range(2)
+      ]
+      for t in threads:
+        t.start()
+      for t in threads:
+        t.join(timeout=10)
     assert not errors, f"cross-loop startup raised: {errors}"
 
   def test_concurrent_stale_cleanup_folds_once(
@@ -10593,24 +10594,27 @@ async def slow_setup(**kwargs):
 
     def run_in_fresh_loop():
       try:
-        with mock.patch.object(plugin, "_lazy_setup", side_effect=slow_setup):
-          barrier.wait(timeout=5)
-          asyncio.run(plugin._ensure_started())
+        barrier.wait(timeout=5)
+        asyncio.run(plugin._ensure_started())
       except BaseException as e:  # noqa: BLE001
         errors.append(e)
 
-    threads = [
-        platform_thread.create_thread(run_in_fresh_loop) for _ in range(2)
-    ]
-    for t in threads:
-      t.start()
-    # Deterministic rendezvous: hold the owner inside setup until BOTH
-    # threads have entered _ensure_started.
-    entered.wait(timeout=5)
-    release.set()
-    for t in threads:
-      t.join(timeout=10)
-      assert not t.is_alive(), "thread failed to terminate"
+    # Patch from this thread only. patch.object swaps a single shared
+    # attribute and is not itself thread safe, so entering it from both
+    # threads raced on _lazy_setup instead of on the code under test.
+    with mock.patch.object(plugin, "_lazy_setup", side_effect=slow_setup):
+      threads = [
+          platform_thread.create_thread(run_in_fresh_loop) for _ in range(2)
+      ]
+      for t in threads:
+        t.start()
+      # Deterministic rendezvous: hold the owner inside setup until BOTH
+      # threads have entered _ensure_started.
+      entered.wait(timeout=5)
+      release.set()
+      for t in threads:
+        t.join(timeout=10)
+        assert not t.is_alive(), "thread failed to terminate"
 
     assert not errors, f"cross-loop startup raised: {errors}"
     assert len(setup_calls) == 1, f"shared setup ran {len(setup_calls)} times"
@@ -10697,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
       )
@@ -10786,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
       )