style: type issues in test folder (#521)

# Description

I noticed there are quite a few type warnings and errors in this repo
when running pyright. I went ahead and fixed some of the type issues in
the `/tests` folder to help clean things up a bit.

Strong typing helps catch bugs early and makes the codebase easier to
maintain in the long run.

I’d be happy to contribute more and tackle additional type problems if
there’s interest. This PR is just a first step to see if it makes sense
to spend more time here and maybe get more involved as a contributor.

Let me know what you think!

---------

Co-authored-by: Holt Skinner <[email protected]>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: 3 people

tests/client/test_auth_middleware.py

@@ -106,10 +106,10 @@ def store():
 
 
 @pytest.mark.asyncio
-async def test_auth_interceptor_skips_when_no_agent_card(store):
-    """
-    Tests that the AuthInterceptor does not modify the request when no AgentCard is provided.
-    """
+async def test_auth_interceptor_skips_when_no_agent_card(
+    store: InMemoryContextCredentialStore,
+) -> None:
+    """Tests that the AuthInterceptor does not modify the request when no AgentCard is provided."""
     request_payload = {'foo': 'bar'}
     http_kwargs = {'fizz': 'buzz'}
     auth_interceptor = AuthInterceptor(credential_service=store)
@@ -126,9 +126,10 @@ async def test_auth_interceptor_skips_when_no_agent_card(store):
 
 
 @pytest.mark.asyncio
-async def test_in_memory_context_credential_store(store):
-    """
-    Verifies that InMemoryContextCredentialStore correctly stores and retrieves
+async def test_in_memory_context_credential_store(
+    store: InMemoryContextCredentialStore,
+) -> None:
+    """Verifies that InMemoryContextCredentialStore correctly stores and retrieves
     credentials based on the session ID in the client context.
     """
     session_id = 'session-id'
@@ -163,11 +164,8 @@ async def test_in_memory_context_credential_store(store):
 
 @pytest.mark.asyncio
 @respx.mock
-async def test_client_with_simple_interceptor():
-    """
-    Ensures that a custom HeaderInterceptor correctly injects a static header
-    into outbound HTTP requests from the A2AClient.
-    """
+async def test_client_with_simple_interceptor() -> None:
+    """Ensures that a custom HeaderInterceptor correctly injects a static header into outbound HTTP requests from the A2AClient."""
     url = 'http://agent.com/rpc'
     interceptor = HeaderInterceptor('X-Test-Header', 'Test-Value-123')
     card = AgentCard(
@@ -196,9 +194,7 @@ async def test_client_with_simple_interceptor():
 
 @dataclass
 class AuthTestCase:
-    """
-    Represents a test scenario for verifying authentication behavior in AuthInterceptor.
-    """
+    """Represents a test scenario for verifying authentication behavior in AuthInterceptor."""
 
     url: str
     """The endpoint URL of the agent to which the request is sent."""
@@ -284,11 +280,10 @@ class AuthTestCase:
     [api_key_test_case, oauth2_test_case, oidc_test_case, bearer_test_case],
 )
 @respx.mock
-async def test_auth_interceptor_variants(test_case, store):
-    """
-    Parametrized test verifying that AuthInterceptor correctly attaches credentials
-    based on the defined security scheme in the AgentCard.
-    """
+async def test_auth_interceptor_variants(
+    test_case: AuthTestCase, store: InMemoryContextCredentialStore
+) -> None:
+    """Parametrized test verifying that AuthInterceptor correctly attaches credentials based on the defined security scheme in the AgentCard."""
     await store.set_credentials(
         test_case.session_id, test_case.scheme_name, test_case.credential
     )
@@ -329,12 +324,9 @@ async def test_auth_interceptor_variants(test_case, store):
 
 @pytest.mark.asyncio
 async def test_auth_interceptor_skips_when_scheme_not_in_security_schemes(
-    store,
-):
-    """
-    Tests that AuthInterceptor skips a scheme if it's listed in security requirements
-    but not defined in security_schemes.
-    """
+    store: InMemoryContextCredentialStore,
+) -> None:
+    """Tests that AuthInterceptor skips a scheme if it's listed in security requirements but not defined in security_schemes."""
     scheme_name = 'missing'
     session_id = 'session-id'
     credential = 'dummy-token'

tests/client/test_base_client.py

@@ -19,12 +19,12 @@
 
 
 @pytest.fixture
-def mock_transport():
+def mock_transport() -> AsyncMock:
     return AsyncMock(spec=ClientTransport)
 
 
 @pytest.fixture
-def sample_agent_card():
+def sample_agent_card() -> AgentCard:
     return AgentCard(
         name='Test Agent',
         description='An agent for testing',
@@ -38,7 +38,7 @@ def sample_agent_card():
 
 
 @pytest.fixture
-def sample_message():
+def sample_message() -> Message:
     return Message(
         role=Role.user,
         message_id='msg-1',
@@ -47,7 +47,9 @@ def sample_message():
 
 
 @pytest.fixture
-def base_client(sample_agent_card, mock_transport):
+def base_client(
+    sample_agent_card: AgentCard, mock_transport: AsyncMock
+) -> BaseClient:
     config = ClientConfig(streaming=True)
     return BaseClient(
         card=sample_agent_card,
@@ -61,7 +63,7 @@ def base_client(sample_agent_card, mock_transport):
 @pytest.mark.asyncio
 async def test_send_message_streaming(
     base_client: BaseClient, mock_transport: MagicMock, sample_message: Message
-):
+) -> None:
     async def create_stream(*args, **kwargs):
         yield Task(
             id='task-123',
@@ -82,7 +84,7 @@ async def create_stream(*args, **kwargs):
 @pytest.mark.asyncio
 async def test_send_message_non_streaming(
     base_client: BaseClient, mock_transport: MagicMock, sample_message: Message
-):
+) -> None:
     base_client._config.streaming = False
     mock_transport.send_message.return_value = Task(
         id='task-456',
@@ -101,7 +103,7 @@ async def test_send_message_non_streaming(
 @pytest.mark.asyncio
 async def test_send_message_non_streaming_agent_capability_false(
     base_client: BaseClient, mock_transport: MagicMock, sample_message: Message
-):
+) -> None:
     base_client._card.capabilities.streaming = False
     mock_transport.send_message.return_value = Task(
         id='task-789',

tests/client/test_client_task_manager.py

@@ -22,12 +22,12 @@
 
 
 @pytest.fixture
-def task_manager():
+def task_manager() -> ClientTaskManager:
     return ClientTaskManager()
 
 
 @pytest.fixture
-def sample_task():
+def sample_task() -> Task:
     return Task(
         id='task123',
         context_id='context456',
@@ -38,29 +38,31 @@ def sample_task():
 
 
 @pytest.fixture
-def sample_message():
+def sample_message() -> Message:
     return Message(
         message_id='msg1',
         role=Role.user,
         parts=[Part(root=TextPart(text='Hello'))],
     )
 
 
-def test_get_task_no_task_id_returns_none(task_manager: ClientTaskManager):
+def test_get_task_no_task_id_returns_none(
+    task_manager: ClientTaskManager,
+) -> None:
     assert task_manager.get_task() is None
 
 
 def test_get_task_or_raise_no_task_raises_error(
     task_manager: ClientTaskManager,
-):
+) -> None:
     with pytest.raises(A2AClientInvalidStateError, match='no current Task'):
         task_manager.get_task_or_raise()
 
 
 @pytest.mark.asyncio
 async def test_save_task_event_with_task(
     task_manager: ClientTaskManager, sample_task: Task
-):
+) -> None:
     await task_manager.save_task_event(sample_task)
     assert task_manager.get_task() == sample_task
     assert task_manager._task_id == sample_task.id
@@ -70,7 +72,7 @@ async def test_save_task_event_with_task(
 @pytest.mark.asyncio
 async def test_save_task_event_with_task_already_set_raises_error(
     task_manager: ClientTaskManager, sample_task: Task
-):
+) -> None:
     await task_manager.save_task_event(sample_task)
     with pytest.raises(
         A2AClientInvalidArgsError,
@@ -82,7 +84,7 @@ async def test_save_task_event_with_task_already_set_raises_error(
 @pytest.mark.asyncio
 async def test_save_task_event_with_status_update(
     task_manager: ClientTaskManager, sample_task: Task, sample_message: Message
-):
+) -> None:
     await task_manager.save_task_event(sample_task)
     status_update = TaskStatusUpdateEvent(
         task_id=sample_task.id,
@@ -98,7 +100,7 @@ async def test_save_task_event_with_status_update(
 @pytest.mark.asyncio
 async def test_save_task_event_with_artifact_update(
     task_manager: ClientTaskManager, sample_task: Task
-):
+) -> None:
     await task_manager.save_task_event(sample_task)
     artifact = Artifact(
         artifact_id='art1', parts=[Part(root=TextPart(text='artifact content'))]
@@ -119,7 +121,7 @@ async def test_save_task_event_with_artifact_update(
 @pytest.mark.asyncio
 async def test_save_task_event_creates_task_if_not_exists(
     task_manager: ClientTaskManager,
-):
+) -> None:
     status_update = TaskStatusUpdateEvent(
         task_id='new_task',
         context_id='new_context',
@@ -135,7 +137,7 @@ async def test_save_task_event_creates_task_if_not_exists(
 @pytest.mark.asyncio
 async def test_process_with_task_event(
     task_manager: ClientTaskManager, sample_task: Task
-):
+) -> None:
     with patch.object(
         task_manager, 'save_task_event', new_callable=AsyncMock
     ) as mock_save:
@@ -144,7 +146,9 @@ async def test_process_with_task_event(
 
 
 @pytest.mark.asyncio
-async def test_process_with_non_task_event(task_manager: ClientTaskManager):
+async def test_process_with_non_task_event(
+    task_manager: ClientTaskManager,
+) -> None:
     with patch.object(
         task_manager, 'save_task_event', new_callable=Mock
     ) as mock_save:
@@ -155,14 +159,14 @@ async def test_process_with_non_task_event(task_manager: ClientTaskManager):
 
 def test_update_with_message(
     task_manager: ClientTaskManager, sample_task: Task, sample_message: Message
-):
+) -> None:
     updated_task = task_manager.update_with_message(sample_message, sample_task)
     assert updated_task.history == [sample_message]
 
 
 def test_update_with_message_moves_status_message(
     task_manager: ClientTaskManager, sample_task: Task, sample_message: Message
-):
+) -> None:
     status_message = Message(
         message_id='status_msg',
         role=Role.agent,