test(a2a): update executor tests to mock runner._get_or_create_session

Author: nickchecan

tests/unittests/a2a/executor/test_a2a_agent_executor.py

@@ -224,19 +224,17 @@ async def mock_run_async(**kwargs):
 
   @pytest.mark.asyncio
   async def test_prepare_session_new_session(self):
-    """Test session preparation when session doesn't exist."""
+    """Test session preparation delegates to runner._get_or_create_session."""
     run_args = AgentRunRequest(
         user_id="test-user",
         session_id=None,
         new_message=Mock(spec=Content),
         run_config=Mock(spec=RunConfig),
     )
 
-    # Mock session service
-    self.mock_runner.session_service.get_session = AsyncMock(return_value=None)
     mock_session = Mock()
     mock_session.id = "new-session-id"
-    self.mock_runner.session_service.create_session = AsyncMock(
+    self.mock_runner._get_or_create_session = AsyncMock(
         return_value=mock_session
     )
 
@@ -245,10 +243,13 @@ async def test_prepare_session_new_session(self):
         self.mock_context, run_args, self.mock_runner
     )
 
-    # Verify session was created
+    # Verify session was returned and run_request updated
     assert result == mock_session
-    assert run_args.session_id is not None
-    self.mock_runner.session_service.create_session.assert_called_once()
+    assert run_args.session_id == "new-session-id"
+    self.mock_runner._get_or_create_session.assert_called_once_with(
+        user_id="test-user",
+        session_id=None,
+    )
 
   @pytest.mark.asyncio
   async def test_prepare_session_existing_session(self):
@@ -260,10 +261,9 @@ async def test_prepare_session_existing_session(self):
         run_config=Mock(spec=RunConfig),
     )
 
-    # Mock session service
     mock_session = Mock()
     mock_session.id = "existing-session"
-    self.mock_runner.session_service.get_session = AsyncMock(
+    self.mock_runner._get_or_create_session = AsyncMock(
         return_value=mock_session
     )
 
@@ -274,7 +274,10 @@ async def test_prepare_session_existing_session(self):
 
     # Verify existing session was returned
     assert result == mock_session
-    self.mock_runner.session_service.create_session.assert_not_called()
+    self.mock_runner._get_or_create_session.assert_called_once_with(
+        user_id="test-user",
+        session_id="existing-session",
+    )
 
   def test_constructor_with_callable_runner(self):
     """Test constructor with callable runner."""

tests/unittests/a2a/executor/test_a2a_agent_executor_impl.py

@@ -97,10 +97,10 @@ async def test_execute_success_new_task(self):
         new_message=Mock(spec=Content),
         run_config=Mock(spec=RunConfig),
     )
-    # Mock session service
+    # Mock _get_or_create_session
     mock_session = Mock()
     mock_session.id = "test-session"
-    self.mock_runner.session_service.get_session = AsyncMock(
+    self.mock_runner._get_or_create_session = AsyncMock(
         return_value=mock_session
     )
 
@@ -200,10 +200,10 @@ async def test_execute_existing_task(self):
         run_config=Mock(spec=RunConfig),
     )
 
-    # Mock session service
+    # Mock _get_or_create_session
     mock_session = Mock()
     mock_session.id = "test-session"
-    self.mock_runner.session_service.get_session = AsyncMock(
+    self.mock_runner._get_or_create_session = AsyncMock(
         return_value=mock_session
     )
 
@@ -616,7 +616,7 @@ async def test_execute_missing_user_input(self, mock_handle_user_input):
     )
     mock_handle_user_input.return_value = missing_event
 
-    self.mock_runner.session_service.get_session = AsyncMock(
+    self.mock_runner._get_or_create_session = AsyncMock(
         return_value=Mock(id="test-session")
     )
     self.mock_request_converter.return_value = AgentRunRequest(
@@ -638,12 +638,10 @@ async def test_execute_missing_user_input(self, mock_handle_user_input):
 
   @pytest.mark.asyncio
   async def test_resolve_session_creates_new_session(self):
-    """Test that _resolve_session creates a new session if it doesn't exist."""
-    self.mock_runner.session_service.get_session = AsyncMock(return_value=None)
-
+    """Test that _resolve_session delegates to runner._get_or_create_session."""
     new_session = Mock()
     new_session.id = "new-session-id"
-    self.mock_runner.session_service.create_session = AsyncMock(
+    self.mock_runner._get_or_create_session = AsyncMock(
         return_value=new_session
     )
 
@@ -656,17 +654,10 @@ async def test_resolve_session_creates_new_session(self):
 
     await self.executor._resolve_session(run_request, self.mock_runner)
 
-    self.mock_runner.session_service.get_session.assert_called_once_with(
-        app_name=self.mock_runner.app_name,
-        user_id="test-user",
-        session_id="old-session-id",
-        config=GetSessionConfig(num_recent_events=0, after_timestamp=None),
-    )
-    self.mock_runner.session_service.create_session.assert_called_once_with(
-        app_name=self.mock_runner.app_name,
+    self.mock_runner._get_or_create_session.assert_called_once_with(
         user_id="test-user",
-        state={},
         session_id="old-session-id",
+        get_session_config=GetSessionConfig(num_recent_events=0, after_timestamp=None),
     )
     assert run_request.session_id == "new-session-id"