Fix more tests

Author: pstephengoogle

tests/server/request_handlers/test_default_request_handler.py

@@ -134,13 +134,13 @@ async def test_on_get_task_not_found():
     params = TaskQueryParams(id='non_existent_task')
 
     from a2a.utils.errors import ServerError  # Local import for ServerError
-
+    context = create_server_call_context()
     with pytest.raises(ServerError) as exc_info:
-        await request_handler.on_get_task(params, create_server_call_context())
+        await request_handler.on_get_task(params, context)
 
     assert isinstance(exc_info.value.error, TaskNotFoundError)
     mock_task_store.get.assert_awaited_once_with(
-        'non_existent_task', create_server_call_context()
+        'non_existent_task', context
     )
 
 
@@ -156,16 +156,16 @@ async def test_on_cancel_task_task_not_found():
     params = TaskIdParams(id='task_not_found_for_cancel')
 
     from a2a.utils.errors import ServerError  # Local import
-
+    context = create_server_call_context()
     with pytest.raises(ServerError) as exc_info:
         await request_handler.on_cancel_task(
-            params, create_server_call_context()
+            params, context
         )
 
     assert isinstance(exc_info.value.error, TaskNotFoundError)
     mock_task_store.get.assert_awaited_once_with(
         'task_not_found_for_cancel',
-        create_server_call_context()
+        context
     )
 
 
@@ -200,18 +200,18 @@ async def test_on_cancel_task_queue_tap_returns_none():
         queue_manager=mock_queue_manager,
     )
 
+    context = create_server_call_context()
     with patch(
         'a2a.server.request_handlers.default_request_handler.ResultAggregator',
         return_value=mock_result_aggregator_instance,
     ):
         params = TaskIdParams(id='tap_none_task')
         result_task = await request_handler.on_cancel_task(
-            params, create_server_call_context()
+            params, context
         )
 
     mock_task_store.get.assert_awaited_once_with(
-        'tap_none_task',
-        create_server_call_context()
+        'tap_none_task', context
     )
     mock_queue_manager.tap.assert_awaited_once_with('tap_none_task')
     # agent_executor.cancel should be called with a new EventQueue if tap returned None
@@ -258,13 +258,14 @@ async def test_on_cancel_task_cancels_running_agent():
     mock_producer_task = AsyncMock(spec=asyncio.Task)
     request_handler._running_agents[task_id] = mock_producer_task
 
+    context = create_server_call_context()
     with patch(
         'a2a.server.request_handlers.default_request_handler.ResultAggregator',
         return_value=mock_result_aggregator_instance,
     ):
         params = TaskIdParams(id=task_id)
         await request_handler.on_cancel_task(
-            params, create_server_call_context()
+            params, context
         )
 
     mock_producer_task.cancel.assert_called_once()
@@ -1279,16 +1280,15 @@ async def test_set_task_push_notification_config_task_not_found():
         ),
     )
     from a2a.utils.errors import ServerError  # Local import
-
+    context = create_server_call_context()
     with pytest.raises(ServerError) as exc_info:
         await request_handler.on_set_task_push_notification_config(
-            params, create_server_call_context()
+            params, context
         )
 
     assert isinstance(exc_info.value.error, TaskNotFoundError)
     mock_task_store.get.assert_awaited_once_with(
-        'non_existent_task',
-        create_server_call_context()
+        'non_existent_task', context
     )
     mock_push_store.set_info.assert_not_awaited()
 
@@ -1326,15 +1326,15 @@ async def test_get_task_push_notification_config_task_not_found():
     params = GetTaskPushNotificationConfigParams(id='non_existent_task')
     from a2a.utils.errors import ServerError  # Local import
 
+    context = create_server_call_context()
     with pytest.raises(ServerError) as exc_info:
         await request_handler.on_get_task_push_notification_config(
-            params, create_server_call_context()
+            params, context
         )
 
     assert isinstance(exc_info.value.error, TaskNotFoundError)
     mock_task_store.get.assert_awaited_once_with(
-        'non_existent_task',
-        create_server_call_context()
+        'non_existent_task', context
     )
     mock_push_store.get_info.assert_not_awaited()
 
@@ -1357,18 +1357,17 @@ async def test_get_task_push_notification_config_info_not_found():
     )
     params = GetTaskPushNotificationConfigParams(id='non_existent_task')
     from a2a.utils.errors import ServerError  # Local import
-
+    context = create_server_call_context()
     with pytest.raises(ServerError) as exc_info:
         await request_handler.on_get_task_push_notification_config(
-            params, create_server_call_context()
+            params, context
         )
 
     assert isinstance(
         exc_info.value.error, InternalError
     )  # Current code raises InternalError
     mock_task_store.get.assert_awaited_once_with(
-        'non_existent_task',
-        create_server_call_context()
+        'non_existent_task', context
     )
     mock_push_store.get_info.assert_awaited_once_with('non_existent_task')
 
@@ -1392,8 +1391,9 @@ async def test_get_task_push_notification_config_info_with_config():
             id='config_id', url='http://1.example.com'
         ),
     )
+    context = create_server_call_context()
     await request_handler.on_set_task_push_notification_config(
-        set_config_params, create_server_call_context()
+        set_config_params, context
     )
 
     params = GetTaskPushNotificationConfigParams(
@@ -1402,7 +1402,7 @@ async def test_get_task_push_notification_config_info_with_config():
 
     result: TaskPushNotificationConfig = (
         await request_handler.on_get_task_push_notification_config(
-            params, create_server_call_context()
+            params, context
         )
     )
 
@@ -1467,18 +1467,17 @@ async def test_on_resubscribe_to_task_task_not_found():
     params = TaskIdParams(id='resub_task_not_found')
 
     from a2a.utils.errors import ServerError  # Local import
-
+    context = create_server_call_context()
     with pytest.raises(ServerError) as exc_info:
         # Need to consume the async generator to trigger the error
         async for _ in request_handler.on_resubscribe_to_task(
-            params, create_server_call_context()
+            params, context
         ):
             pass
 
     assert isinstance(exc_info.value.error, TaskNotFoundError)
     mock_task_store.get.assert_awaited_once_with(
-        'resub_task_not_found',
-        create_server_call_context()
+        'resub_task_not_found', context
     )
 
 
@@ -1501,18 +1500,18 @@ async def test_on_resubscribe_to_task_queue_not_found():
 
     from a2a.utils.errors import ServerError  # Local import
 
+    context = create_server_call_context()
     with pytest.raises(ServerError) as exc_info:
         async for _ in request_handler.on_resubscribe_to_task(
-            params, create_server_call_context()
+            params, context
         ):
             pass
 
     assert isinstance(
         exc_info.value.error, TaskNotFoundError
     )  # Should be TaskNotFoundError as per spec
     mock_task_store.get.assert_awaited_once_with(
-        'resub_queue_not_found',
-        create_server_call_context()
+        'resub_queue_not_found', context
     )
     mock_queue_manager.tap.assert_awaited_once_with('resub_queue_not_found')
 
@@ -1586,16 +1585,15 @@ async def test_list_task_push_notification_config_task_not_found():
     )
     params = ListTaskPushNotificationConfigParams(id='non_existent_task')
     from a2a.utils.errors import ServerError  # Local import
-
+    context = create_server_call_context()
     with pytest.raises(ServerError) as exc_info:
         await request_handler.on_list_task_push_notification_config(
-            params, create_server_call_context()
+            params, context
         )
 
     assert isinstance(exc_info.value.error, TaskNotFoundError)
     mock_task_store.get.assert_awaited_once_with(
-        'non_existent_task',
-        create_server_call_context()
+        'non_existent_task', context
     )
     mock_push_store.get_info.assert_not_awaited()
 
@@ -1749,16 +1747,15 @@ async def test_delete_task_push_notification_config_task_not_found():
         id='non_existent_task', push_notification_config_id='config1'
     )
     from a2a.utils.errors import ServerError  # Local import
-
+    context = create_server_call_context()
     with pytest.raises(ServerError) as exc_info:
         await request_handler.on_delete_task_push_notification_config(
-            params, create_server_call_context()
+            params, context
         )
 
     assert isinstance(exc_info.value.error, TaskNotFoundError)
     mock_task_store.get.assert_awaited_once_with(
-        'non_existent_task',
-        create_server_call_context()
+        'non_existent_task', context
     )
     mock_push_store.get_info.assert_not_awaited()
 
@@ -2003,10 +2000,10 @@ async def test_on_resubscribe_to_task_in_terminal_state(terminal_state):
     params = TaskIdParams(id=task_id)
 
     from a2a.utils.errors import ServerError
-
+    context = create_server_call_context()
     with pytest.raises(ServerError) as exc_info:
         async for _ in request_handler.on_resubscribe_to_task(
-            params, create_server_call_context()
+            params, context
         ):
             pass  # pragma: no cover
 
@@ -2017,7 +2014,7 @@ async def test_on_resubscribe_to_task_in_terminal_state(terminal_state):
         in exc_info.value.error.message
     )
     mock_task_store.get.assert_awaited_once_with(
-        task_id, create_server_call_context()
+        task_id, context
     )