fix: type issues

Author: HelloJowet

src/a2a/server/tasks/__init__.py

@@ -2,21 +2,17 @@
 
 import logging
 
-from a2a.server.tasks.base_push_notification_sender import (
+from a2a.server.tasks import (
     BasePushNotificationSender,
-)
-from a2a.server.tasks.inmemory_push_notification_config_store import (
     InMemoryPushNotificationConfigStore,
-)
-from a2a.server.tasks.inmemory_task_store import InMemoryTaskStore
-from a2a.server.tasks.push_notification_config_store import (
+    InMemoryTaskStore,
     PushNotificationConfigStore,
+    PushNotificationSender,
+    ResultAggregator,
+    TaskManager,
+    TaskStore,
+    TaskUpdater,
 )
-from a2a.server.tasks.push_notification_sender import PushNotificationSender
-from a2a.server.tasks.result_aggregator import ResultAggregator
-from a2a.server.tasks.task_manager import TaskManager
-from a2a.server.tasks.task_store import TaskStore
-from a2a.server.tasks.task_updater import TaskUpdater
 
 
 logger = logging.getLogger(__name__)

tests/client/test_auth_middleware.py

@@ -17,7 +17,6 @@
     ClientFactory,
     InMemoryContextCredentialStore,
 )
-from a2a.client.auth.credentials import InMemoryContextCredentialStore
 from a2a.types import (
     APIKeySecurityScheme,
     AgentCapabilities,

tests/client/test_base_client.py

@@ -47,7 +47,9 @@ def sample_message() -> Message:
 
 
 @pytest.fixture
-def base_client(sample_agent_card: AgentCard, mock_transport: AsyncMock):
+def base_client(
+    sample_agent_card: AgentCard, mock_transport: AsyncMock
+) -> BaseClient:
     config = ClientConfig(streaming=True)
     return BaseClient(
         card=sample_agent_card,

tests/server/agent_execution/test_context.py

@@ -106,7 +106,7 @@ def test_get_user_input_no_params(self):
         context = RequestContext()
         assert context.get_user_input() == ''
 
-    def test_attach_related_task(self, mock_task: Mock):
+    def test_attach_related_task(self, mock_task: Mock) -> None:
         """Test attach_related_task adds a task to related_tasks."""
         context = RequestContext()
         assert len(context.related_tasks) == 0
@@ -257,7 +257,7 @@ def test_metadata_property_with_content(self, mock_params: Mock) -> None:
         assert context.metadata == {'key': 'value'}
 
     def test_init_with_existing_ids_in_message(
-        self, mock_message, mock_params
+        self, mock_message: Mock, mock_params: Mock
     ) -> None:
         """Test initialization with existing IDs in the message."""
         mock_message.task_id = 'existing-task-id'
@@ -270,7 +270,7 @@ def test_init_with_existing_ids_in_message(
         # No new UUIDs should be generated
 
     def test_init_with_task_id_and_existing_task_id_match(
-        self, mock_params: Mock, mock_task
+        self, mock_params: Mock, mock_task: Mock
     ) -> None:
         """Test initialization succeeds when task_id matches task.id."""
         mock_params.message.task_id = mock_task.id

tests/server/agent_execution/test_simple_request_context_builder.py

@@ -26,23 +26,25 @@
 
 # Helper to create a simple message
 def create_sample_message(
-    content='test message',
-    msg_id='msg1',
-    role=Role.user,
-    reference_task_ids=None,
-):
+    content: str = 'test message',
+    msg_id: str = 'msg1',
+    role: Role = Role.user,
+    reference_task_ids: list[str] | None = None,
+) -> Message:
     return Message(
         message_id=msg_id,
         role=role,
         parts=[Part(root=TextPart(text=content))],
-        referenceTaskIds=reference_task_ids if reference_task_ids else [],
+        reference_task_ids=reference_task_ids if reference_task_ids else [],
     )
 
 
 # Helper to create a simple task
 def create_sample_task(
-    task_id='task1', status_state=TaskState.submitted, context_id='ctx1'
-):
+    task_id: str = 'task1',
+    status_state: TaskState = TaskState.submitted,
+    context_id: str = 'ctx1',
+) -> Task:
     return Task(
         id=task_id,
         context_id=context_id,

tests/server/events/test_inmemory_queue_manager.py

@@ -28,7 +28,7 @@ def event_queue(self) -> MagicMock:
         return queue
 
     @pytest.mark.asyncio
-    async def test_init(self, queue_manager: InMemoryQueueManager):
+    async def test_init(self, queue_manager: InMemoryQueueManager) -> None:
         """Test that the InMemoryQueueManager initializes with empty task queue and a lock."""
         assert queue_manager._task_queue == {}
         assert isinstance(queue_manager._lock, asyncio.Lock)

tests/server/request_handlers/test_response_helpers.py

@@ -93,7 +93,9 @@ def test_build_error_response_with_request_id_none(self):
         self.assertIsInstance(response_wrapper.root, JSONRPCErrorResponse)
         self.assertIsNone(response_wrapper.root.id)
 
-    def _create_sample_task(self, task_id='task123', context_id='ctx456'):
+    def _create_sample_task(
+        self, task_id: str = 'task123', context_id: str = 'ctx456'
+    ) -> Task:
         return Task(
             id=task_id,
             context_id=context_id,

tests/server/tasks/test_inmemory_push_notifications.py

@@ -19,7 +19,7 @@
 
 def create_sample_task(
     task_id: str = 'task123', status_state: TaskState = TaskState.completed
-):
+) -> Task:
     return Task(
         id=task_id,
         context_id='ctx456',
@@ -31,7 +31,7 @@ def create_sample_push_config(
     url: str = 'http://example.com/callback',
     config_id: str = 'cfg1',
     token: str | None = None,
-):
+) -> PushNotificationConfig:
     return PushNotificationConfig(id=config_id, url=url, token=token)
 
 

tests/server/tasks/test_push_notification_sender.py

@@ -17,7 +17,7 @@
 
 def create_sample_task(
     task_id: str = 'task123', status_state: TaskState = TaskState.completed
-):
+) -> Task:
     return Task(
         id=task_id,
         context_id='ctx456',
@@ -29,7 +29,7 @@ def create_sample_push_config(
     url: str = 'http://example.com/callback',
     config_id: str = 'cfg1',
     token: str | None = None,
-):
+) -> PushNotificationConfig:
     return PushNotificationConfig(id=config_id, url=url, token=token)
 
 

tests/server/tasks/test_result_aggregator.py

@@ -1,4 +1,5 @@
 import asyncio
+from typing_extensions import override
 import unittest
 
 from collections.abc import AsyncIterator
@@ -22,7 +23,7 @@
 # Helper to create a simple message
 def create_sample_message(
     content: str = 'test message', msg_id: str = 'msg1', role: Role = Role.user
-):
+) -> Message:
     return Message(
         message_id=msg_id,
         role=role,
@@ -35,7 +36,7 @@ def create_sample_task(
     task_id: str = 'task1',
     status_state: TaskState = TaskState.submitted,
     context_id: str = 'ctx1',
-):
+) -> Task:
     return Task(
         id=task_id,
         context_id=context_id,
@@ -45,8 +46,10 @@ def create_sample_task(
 
 # Helper to create a TaskStatusUpdateEvent
 def create_sample_status_update(
-    task_id='task1', status_state=TaskState.working, context_id='ctx1'
-):
+    task_id: str = 'task1',
+    status_state: TaskState = TaskState.working,
+    context_id: str = 'ctx1',
+) -> TaskStatusUpdateEvent:
     return TaskStatusUpdateEvent(
         task_id=task_id,
         context_id=context_id,
@@ -56,6 +59,7 @@ def create_sample_status_update(
 
 
 class TestResultAggregator(unittest.IsolatedAsyncioTestCase):
+    @override
     def setUp(self):
         self.mock_task_manager = AsyncMock(spec=TaskManager)
         self.mock_event_consumer = AsyncMock(spec=EventConsumer)