Update more tests

Author: pstephengoogle

tests/server/request_handlers/test_default_request_handler.py

@@ -139,7 +139,9 @@ async def test_on_get_task_not_found():
         await request_handler.on_get_task(params, create_server_call_context())
 
     assert isinstance(exc_info.value.error, TaskNotFoundError)
-    mock_task_store.get.assert_awaited_once_with('non_existent_task')
+    mock_task_store.get.assert_awaited_once_with(
+        'non_existent_task', create_server_call_context()
+    )
 
 
 @pytest.mark.asyncio
@@ -161,7 +163,10 @@ async def test_on_cancel_task_task_not_found():
         )
 
     assert isinstance(exc_info.value.error, TaskNotFoundError)
-    mock_task_store.get.assert_awaited_once_with('task_not_found_for_cancel')
+    mock_task_store.get.assert_awaited_once_with(
+        'task_not_found_for_cancel',
+        create_server_call_context()
+    )
 
 
 @pytest.mark.asyncio
@@ -204,7 +209,10 @@ async def test_on_cancel_task_queue_tap_returns_none():
             params, create_server_call_context()
         )
 
-    mock_task_store.get.assert_awaited_once_with('tap_none_task')
+    mock_task_store.get.assert_awaited_once_with(
+        'tap_none_task',
+        create_server_call_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
     mock_agent_executor.cancel.assert_awaited_once()
@@ -1278,7 +1286,10 @@ async def test_set_task_push_notification_config_task_not_found():
         )
 
     assert isinstance(exc_info.value.error, TaskNotFoundError)
-    mock_task_store.get.assert_awaited_once_with('non_existent_task')
+    mock_task_store.get.assert_awaited_once_with(
+        'non_existent_task',
+        create_server_call_context()
+    )
     mock_push_store.set_info.assert_not_awaited()
 
 
@@ -1321,7 +1332,10 @@ async def test_get_task_push_notification_config_task_not_found():
         )
 
     assert isinstance(exc_info.value.error, TaskNotFoundError)
-    mock_task_store.get.assert_awaited_once_with('non_existent_task')
+    mock_task_store.get.assert_awaited_once_with(
+        'non_existent_task',
+        create_server_call_context()
+    )
     mock_push_store.get_info.assert_not_awaited()
 
 
@@ -1352,7 +1366,10 @@ async def test_get_task_push_notification_config_info_not_found():
     assert isinstance(
         exc_info.value.error, InternalError
     )  # Current code raises InternalError
-    mock_task_store.get.assert_awaited_once_with('non_existent_task')
+    mock_task_store.get.assert_awaited_once_with(
+        'non_existent_task',
+        create_server_call_context()
+    )
     mock_push_store.get_info.assert_awaited_once_with('non_existent_task')
 
 
@@ -1459,7 +1476,10 @@ async def test_on_resubscribe_to_task_task_not_found():
             pass
 
     assert isinstance(exc_info.value.error, TaskNotFoundError)
-    mock_task_store.get.assert_awaited_once_with('resub_task_not_found')
+    mock_task_store.get.assert_awaited_once_with(
+        'resub_task_not_found',
+        create_server_call_context()
+    )
 
 
 @pytest.mark.asyncio
@@ -1490,7 +1510,10 @@ async def test_on_resubscribe_to_task_queue_not_found():
     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')
+    mock_task_store.get.assert_awaited_once_with(
+        'resub_queue_not_found',
+        create_server_call_context()
+    )
     mock_queue_manager.tap.assert_awaited_once_with('resub_queue_not_found')
 
 
@@ -1570,7 +1593,10 @@ async def test_list_task_push_notification_config_task_not_found():
         )
 
     assert isinstance(exc_info.value.error, TaskNotFoundError)
-    mock_task_store.get.assert_awaited_once_with('non_existent_task')
+    mock_task_store.get.assert_awaited_once_with(
+        'non_existent_task',
+        create_server_call_context()
+    )
     mock_push_store.get_info.assert_not_awaited()
 
 
@@ -1730,7 +1756,10 @@ async def test_delete_task_push_notification_config_task_not_found():
         )
 
     assert isinstance(exc_info.value.error, TaskNotFoundError)
-    mock_task_store.get.assert_awaited_once_with('non_existent_task')
+    mock_task_store.get.assert_awaited_once_with(
+        'non_existent_task',
+        create_server_call_context()
+    )
     mock_push_store.get_info.assert_not_awaited()
 
 
@@ -1987,7 +2016,9 @@ async def test_on_resubscribe_to_task_in_terminal_state(terminal_state):
         f'Task {task_id} is in terminal state: {terminal_state.value}'
         in exc_info.value.error.message
     )
-    mock_task_store.get.assert_awaited_once_with(task_id)
+    mock_task_store.get.assert_awaited_once_with(
+        task_id, create_server_call_context()
+    )
 
 
 @pytest.mark.asyncio

tests/server/tasks/test_task_manager.py

@@ -116,7 +116,7 @@ async def test_save_task_event_status_update(
     await task_manager.save_task_event(event)
     updated_task = initial_task
     updated_task.status = new_status
-    mock_task_store.save.assert_called_once_with(updated_task)
+    mock_task_store.save.assert_called_once_with(updated_task, None)
 
 
 @pytest.mark.asyncio
@@ -139,7 +139,7 @@ async def test_save_task_event_artifact_update(
     await task_manager.save_task_event(event)
     updated_task = initial_task
     updated_task.artifacts = [new_artifact]
-    mock_task_store.save.assert_called_once_with(updated_task)
+    mock_task_store.save.assert_called_once_with(updated_task, None)
 
 
 @pytest.mark.asyncio