fix: wait for in-flight BigQuery writes

BatchProcessor.flush() checked Queue.empty(), which only reports whether an item is still waiting, so it could return after a writer dequeued an item but before that write finished. This waits on the queue's unfinished-task count with Queue.join() so flush() blocks until in-flight writes complete.

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

Author: GWeale

src/google/adk/plugins/bigquery_agent_analytics_plugin.py

@@ -1484,10 +1484,9 @@ def __init__(
     }
 
   async def flush(self) -> None:
-    """Flushes the queue by waiting for it to be empty."""
-    if self._queue.empty():
-      return
-    # Wait for all items in the queue to be processed
+    """Flushes the queue, blocking until in-flight writes complete."""
+    # empty() turns true as soon as an item is dequeued, before its write
+    # finishes; join() waits for the unfinished-task count to reach zero.
     await self._queue.join()
 
   async def start(self) -> None:

tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py

@@ -8547,6 +8547,19 @@ def _stub_arrow_prep(self, bp):
     fake_batch.serialize.return_value.to_pybytes.return_value = b"batch"
     bp._prepare_arrow_batch = mock.MagicMock(return_value=fake_batch)
 
+  @pytest.mark.asyncio
+  async def test_flush_waits_for_dequeued_write(self, dummy_arrow_schema):
+    bp = self._make_processor(dummy_arrow_schema)
+    await bp.append({"event": 0})
+    await bp._queue.get()
+
+    flush_task = asyncio.create_task(bp.flush())
+    await asyncio.sleep(0)
+
+    assert not flush_task.done()
+    bp._queue.task_done()
+    await flush_task
+
   @pytest.mark.asyncio
   async def test_queue_full_drops_are_counted(self, dummy_arrow_schema):
     # Writer is not started, so a size-1 queue fills after one append and the