More tests

Author: pstephengoogle

tests/server/request_handlers/test_jsonrpc_handler.py

@@ -113,7 +113,7 @@ async def test_on_get_task_success(self) -> None:
         )
         self.assertIsInstance(response.root, GetTaskSuccessResponse)
         assert response.root.result == mock_task  # type: ignore
-        mock_task_store.get.assert_called_once_with(task_id)
+        mock_task_store.get.assert_called_once_with(task_id, unittest.mock.ANY)
 
     async def test_on_get_task_not_found(self) -> None:
         mock_agent_executor = AsyncMock(spec=AgentExecutor)
@@ -206,7 +206,9 @@ async def test_on_cancel_task_not_found(self) -> None:
         response = await handler.on_cancel_task(request)
         self.assertIsInstance(response.root, JSONRPCErrorResponse)
         assert response.root.error == TaskNotFoundError()  # type: ignore
-        mock_task_store.get.assert_called_once_with('nonexistent_id')
+        mock_task_store.get.assert_called_once_with(
+            'nonexistent_id', unittest.mock.ANY
+        )
         mock_agent_executor.cancel.assert_not_called()
 
     @patch(