feat: Improve asyncio loop handling and test cleanup

This CL enhances asyncio event loop management and test isolation. - **BigQuery Analytics Plugin:** Ensure the asyncio event loop is consistently closed within the BigQuery analytics plugin. This prevents potential resource leaks. Add checks to handle potential deadlocks in Python 3.13+ when creating loops during interpreter shutdown. - **Test Thread Pool Cleanup:** Introduce a pytest fixture (`cleanup_thread_pools`) to automatically shut down and clear all tool-related thread pools after each test run in `test_functions_thread_pool.py`. This improves test isolation and prevents order-dependent test failures. - **Streaming Test Loop Restoration:** Refactor event loop handling in `test_streaming.py`. A new `_run_with_loop` method is introduced in the custom test runners to create a temporary event loop for each test execution, run the coroutine, and crucially, restore the original event loop afterwards. This prevents tests from interfering with each other's loop state. - **Resource Closure:** Ensure services are closed properly in tests by adding `await service.close()` in `test_service_factory.py` and using `async with session_service` in `test_session_service.py`. PiperOrigin-RevId: 863305565
2026-07-09 18:19:28 -07:00 · 2026-01-30 10:47:09 -08:00
parent 585ebfdac7
commit 00aba2d884
6 changed files with 2334 additions and 2241 deletions
@@ -509,60 +509,61 @@ async def test_append_event_to_stale_session():
      service_type=SessionServiceType.DATABASE
  )

-  app_name = 'my_app'
-  user_id = 'user'
-  current_time = datetime.now().astimezone(timezone.utc).timestamp()
+  async with session_service:
+    app_name = 'my_app'
+    user_id = 'user'
+    current_time = datetime.now().astimezone(timezone.utc).timestamp()

-  original_session = await session_service.create_session(
-      app_name=app_name, user_id=user_id
-  )
-  event1 = Event(
-      invocation_id='inv1',
-      author='user',
-      timestamp=current_time + 1,
-      actions=EventActions(state_delta={'sk1': 'v1'}),
-  )
-  await session_service.append_event(original_session, event1)
+    original_session = await session_service.create_session(
+        app_name=app_name, user_id=user_id
+    )
+    event1 = Event(
+        invocation_id='inv1',
+        author='user',
+        timestamp=current_time + 1,
+        actions=EventActions(state_delta={'sk1': 'v1'}),
+    )
+    await session_service.append_event(original_session, event1)

-  updated_session = await session_service.get_session(
-      app_name=app_name, user_id=user_id, session_id=original_session.id
-  )
-  event2 = Event(
-      invocation_id='inv2',
-      author='user',
-      timestamp=current_time + 2,
-      actions=EventActions(state_delta={'sk2': 'v2'}),
-  )
-  await session_service.append_event(updated_session, event2)
+    updated_session = await session_service.get_session(
+        app_name=app_name, user_id=user_id, session_id=original_session.id
+    )
+    event2 = Event(
+        invocation_id='inv2',
+        author='user',
+        timestamp=current_time + 2,
+        actions=EventActions(state_delta={'sk2': 'v2'}),
+    )
+    await session_service.append_event(updated_session, event2)

-  # original_session is now stale
-  assert original_session.last_update_time < updated_session.last_update_time
-  assert len(original_session.events) == 1
-  assert 'sk2' not in original_session.state
+    # original_session is now stale
+    assert original_session.last_update_time < updated_session.last_update_time
+    assert len(original_session.events) == 1
+    assert 'sk2' not in original_session.state

-  # Appending another event to stale original_session
-  event3 = Event(
-      invocation_id='inv3',
-      author='user',
-      timestamp=current_time + 3,
-      actions=EventActions(state_delta={'sk3': 'v3'}),
-  )
-  await session_service.append_event(original_session, event3)
+    # Appending another event to stale original_session
+    event3 = Event(
+        invocation_id='inv3',
+        author='user',
+        timestamp=current_time + 3,
+        actions=EventActions(state_delta={'sk3': 'v3'}),
+    )
+    await session_service.append_event(original_session, event3)

-  # If we fetch session from DB, it should contain all 3 events and all state
-  # changes.
-  session_final = await session_service.get_session(
-      app_name=app_name, user_id=user_id, session_id=original_session.id
-  )
-  assert len(session_final.events) == 3
-  assert session_final.state.get('sk1') == 'v1'
-  assert session_final.state.get('sk2') == 'v2'
-  assert session_final.state.get('sk3') == 'v3'
-  assert [e.invocation_id for e in session_final.events] == [
-      'inv1',
-      'inv2',
-      'inv3',
-  ]
+    # If we fetch session from DB, it should contain all 3 events and all state
+    # changes.
+    session_final = await session_service.get_session(
+        app_name=app_name, user_id=user_id, session_id=original_session.id
+    )
+    assert len(session_final.events) == 3
+    assert session_final.state.get('sk1') == 'v1'
+    assert session_final.state.get('sk2') == 'v2'
+    assert session_final.state.get('sk3') == 'v3'
+    assert [e.invocation_id for e in session_final.events] == [
+        'inv1',
+        'inv2',
+        'inv3',
+    ]


@pytest.mark.asyncio