core(tests): fix broken tests

Author: thiagosalvatore

pyproject.toml

@@ -149,3 +149,4 @@ line_length = 200
 [tool.pytest.ini_options]
 testpaths = ["tests"]
 asyncio_default_fixture_loop_scope = "function"
+asyncio_mode = "auto" # or "strict"

tests/unittests/flows/llm_flows/test_functions_request_euc.py

@@ -309,7 +309,7 @@ def call_external_api2(tool_context: ToolContext) -> int:
   )
   runner = utils.InMemoryRunner(agent)
   await runner.run('test')
-  request_euc_function_call_event = runner.session.events[-3]
+  request_euc_function_call_event = (await runner.session).events[-3]
   function_response1 = types.FunctionResponse(
       name=request_euc_function_call_event.content.parts[0].function_call.name,
       response=auth_response1.model_dump(),

tests/unittests/flows/llm_flows/test_functions_simple.py

@@ -231,7 +231,7 @@ def update_state(tool_context: ToolContext):
   agent = Agent(name='root_agent', model=mock_model, tools=[update_state])
   runner = utils.InMemoryRunner(agent)
   await runner.run('test')
-  assert runner.session.state['x'] == 1
+  assert (await runner.session).state['x'] == 1
 
 
 async def test_function_call_id():

tests/unittests/sessions/test_session_service.py

@@ -63,7 +63,7 @@ async def test_create_get_session(service_type):
   assert session.id
   assert session.state == state
   assert (
-      session_service.get_session(
+      await session_service.get_session(
           app_name=app_name, user_id=user_id, session_id=session.id
       )
       == session

tests/unittests/sessions/test_vertex_ai_session_service.py

@@ -124,7 +124,7 @@ def __init__(self) -> None:
     this.session_dict: dict[str, Any] = {}
     this.event_dict: dict[str, list[Any]] = {}
 
-  def request(self, http_method: str, path: str, request_dict: dict[str, Any]):
+  async def async_request(self, http_method: str, path: str, request_dict: dict[str, Any]):
     """Mocks the API Client request method."""
     if http_method == 'GET':
       if re.match(SESSION_REGEX, path):

tests/unittests/tools/test_agent_tool.py

@@ -103,13 +103,13 @@ async def test_update_state():
   )
 
   runner = utils.InMemoryRunner(root_agent)
-  runner.session.state['state_1'] = 'state1_value'
+  (await runner.session).state['state_1'] = 'state1_value'
 
   await runner.run('test1')
   assert (
       'input: changed_value' in mock_model.requests[1].config.system_instruction
   )
-  assert runner.session.state['state_1'] == 'changed_value'
+  assert (await runner.session).state['state_1'] == 'changed_value'
 
 
 @mark.parametrize(
@@ -151,7 +151,7 @@ class CustomOutput(BaseModel):
   )
 
   runner = utils.InMemoryRunner(root_agent)
-  runner.session.state['state_1'] = 'state1_value'
+  (await runner.session).state['state_1'] = 'state1_value'
 
   assert utils.simplify_events(await runner.run('test1')) == [
       ('root_agent', function_call_custom),

tests/unittests/utils.py

@@ -172,13 +172,15 @@ def __init__(
         session_service=InMemorySessionService(),
         memory_service=InMemoryMemoryService(),
     )
+    self.session_id = None
 
 
   @property
   async def session(self) -> Session:
-    self.session_id = (await self.runner.session_service.create_session(
-        app_name='test_app', user_id='test_user'
-    )).id
+    if not self.session_id:
+        self.session_id = (await self.runner.session_service.create_session(
+            app_name='test_app', user_id='test_user'
+        )).id
     return await self.runner.session_service.get_session(
         app_name='test_app', user_id='test_user', session_id=self.session_id
     )