test: improve coverage

Signed-off-by: Shingo OKAWA <[email protected]>

Author: ognis1205

src/a2a/utils/proto_utils.py

@@ -14,8 +14,10 @@
 
 
 # Regexp patterns for matching
-_TASK_NAME_MATCH = r'tasks/(\w+)'
-_TASK_PUSH_CONFIG_NAME_MATCH = r'tasks/(\w+)/pushNotificationConfigs/(\w+)'
+_TASK_NAME_MATCH = r'tasks/([a-zA-Z0-9_.-]+)'
+_TASK_PUSH_CONFIG_NAME_MATCH = (
+    r'tasks/([a-zA-Z0-9_.-]+)/pushNotificationConfigs/([a-zA-Z0-9_.-]+)'
+)
 
 
 class ToProto:
@@ -631,6 +633,7 @@ def task_push_notification_config(
         request: a2a_pb2.CreateTaskPushNotificationConfigRequest,
     ) -> types.TaskPushNotificationConfig:
         m = re.match(_TASK_NAME_MATCH, request.parent)
+        print(m)
         if not m:
             raise ServerError(
                 error=types.InvalidParamsError(

tests/client/test_grpc_client.py

@@ -26,6 +26,7 @@
     TextPart,
 )
 from a2a.utils import proto_utils
+from a2a.utils.errors import ServerError
 
 
 # Fixtures
@@ -38,7 +39,7 @@ def mock_grpc_stub() -> AsyncMock:
     stub.GetTask = AsyncMock()
     stub.CancelTask = AsyncMock()
     stub.CreateTaskPushNotificationConfig = AsyncMock()
-    stub.GetTaskPushNotification = AsyncMock()
+    stub.GetTaskPushNotificationConfig = AsyncMock()
     return stub
 
 
@@ -292,3 +293,109 @@ async def test_cancel_task(
         a2a_pb2.CancelTaskRequest(name=f'tasks/{sample_task.id}')
     )
     assert response.status.state == TaskState.canceled
+
+
+@pytest.mark.asyncio
+async def test_set_task_callback_with_valid_task(
+    grpc_client: A2AGrpcClient,
+    mock_grpc_stub: AsyncMock,
+    sample_task_push_notification_config: TaskPushNotificationConfig,
+):
+    """Test setting a task push notification config with a valid task id."""
+    mock_grpc_stub.CreateTaskPushNotificationConfig.return_value = (
+        a2a_pb2.CreateTaskPushNotificationConfigRequest(
+            parent='tasks/task-1',
+            config_id='config-1',
+            config=proto_utils.ToProto.task_push_notification_config(
+                sample_task_push_notification_config
+            ),
+        )
+    )
+
+    response = await grpc_client.set_task_callback(
+        sample_task_push_notification_config
+    )
+
+    mock_grpc_stub.CreateTaskPushNotificationConfig.assert_awaited_once_with(
+        a2a_pb2.CreateTaskPushNotificationConfigRequest(
+            config=proto_utils.ToProto.task_push_notification_config(
+                sample_task_push_notification_config
+            ),
+        )
+    )
+    assert response.taskId == 'task-1'
+
+
+@pytest.mark.asyncio
+async def test_set_task_callback_with_invalid_task(
+    grpc_client: A2AGrpcClient,
+    mock_grpc_stub: AsyncMock,
+    sample_task_push_notification_config: TaskPushNotificationConfig,
+):
+    """Test setting a task push notification config with a invalid task id."""
+    mock_grpc_stub.CreateTaskPushNotificationConfig.return_value = (
+        a2a_pb2.CreateTaskPushNotificationConfigRequest(
+            parent='invalid-path-to-tasks/task-1',
+            config_id='config-1',
+            config=proto_utils.ToProto.task_push_notification_config(
+                sample_task_push_notification_config
+            ),
+        )
+    )
+
+    with pytest.raises(ServerError) as exc_info:
+        await grpc_client.set_task_callback(
+            sample_task_push_notification_config
+        )
+    assert 'No task for' in exc_info.value.error.message
+
+
+@pytest.mark.asyncio
+async def test_get_task_callback_with_valid_task(
+    grpc_client: A2AGrpcClient,
+    mock_grpc_stub: AsyncMock,
+    sample_task_push_notification_config: TaskPushNotificationConfig,
+):
+    """Test retrieving a task push notification config with a valid task id."""
+    mock_grpc_stub.GetTaskPushNotificationConfig.return_value = (
+        a2a_pb2.CreateTaskPushNotificationConfigRequest(
+            parent='tasks/task-1',
+            config_id='config-1',
+            config=proto_utils.ToProto.task_push_notification_config(
+                sample_task_push_notification_config
+            ),
+        )
+    )
+    params = TaskIdParams(id=sample_task_push_notification_config.taskId)
+
+    response = await grpc_client.get_task_callback(params)
+
+    mock_grpc_stub.GetTaskPushNotificationConfig.assert_awaited_once_with(
+        a2a_pb2.GetTaskPushNotificationConfigRequest(
+            name=f'tasks/{params.id}/pushNotification/undefined',
+        )
+    )
+    assert response.taskId == 'task-1'
+
+
+@pytest.mark.asyncio
+async def test_get_task_callback_with_invalid_task(
+    grpc_client: A2AGrpcClient,
+    mock_grpc_stub: AsyncMock,
+    sample_task_push_notification_config: TaskPushNotificationConfig,
+):
+    """Test retrieving a task push notification config with a invalid task id."""
+    mock_grpc_stub.GetTaskPushNotificationConfig.return_value = (
+        a2a_pb2.CreateTaskPushNotificationConfigRequest(
+            parent='invalid-path-to-tasks/task-1',
+            config_id='config-1',
+            config=proto_utils.ToProto.task_push_notification_config(
+                sample_task_push_notification_config
+            ),
+        )
+    )
+    params = TaskIdParams(id=sample_task_push_notification_config.taskId)
+
+    with pytest.raises(ServerError) as exc_info:
+        await grpc_client.get_task_callback(params)
+    assert 'No task for' in exc_info.value.error.message