use gather for async dispatch

Author: kthota-g

src/a2a/server/request_handlers/default_request_handler.py

@@ -411,7 +411,11 @@ async def on_get_task_push_notification_config(
             params.id
         )
         if not push_notification_config or not push_notification_config[0]:
-            raise ServerError(error=InternalError())
+            raise ServerError(
+                error=InternalError(
+                    message='Push notification config not found'
+                )
+            )
 
         return TaskPushNotificationConfig(
             taskId=params.id, pushNotificationConfig=push_notification_config[0]

src/a2a/server/tasks/base_push_notification_sender.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 
 import httpx
@@ -15,7 +16,11 @@
 class BasePushNotificationSender(PushNotificationSender):
     """Base implementation of PushNotificationSender interface."""
 
-    def __init__(self, httpx_client: httpx.AsyncClient, config_store: PushNotificationConfigStore) -> None:
+    def __init__(
+        self,
+        httpx_client: httpx.AsyncClient,
+        config_store: PushNotificationConfigStore,
+    ) -> None:
         """Initializes the BasePushNotificationSender.
 
         Args:
@@ -31,16 +36,32 @@ async def send_notification(self, task: Task) -> None:
         if not push_configs:
             return
 
-        for push_info in push_configs:
-            await self._dispatch_notification(task, push_info)
+        awaitables = [
+            self._dispatch_notification(task, push_info)
+            for push_info in push_configs
+        ]
+        results = await asyncio.gather(*awaitables)
 
-    async def _dispatch_notification(self, task: Task, push_info: PushNotificationConfig) -> None:
+        if not all(results):
+            logger.warning(
+                f'Some push notifications failed to send for task_id={task.id}'
+            )
+
+    async def _dispatch_notification(
+        self, task: Task, push_info: PushNotificationConfig
+    ) -> bool:
         url = push_info.url
         try:
             response = await self._client.post(
                 url, json=task.model_dump(mode='json', exclude_none=True)
             )
             response.raise_for_status()
-            logger.info(f'Push-notification sent for URL: {url}')
+            logger.info(
+                f'Push-notification sent for task_id={task.id} to URL: {url}'
+            )
+            return True
         except Exception as e:
-            logger.error(f'Error sending push-notification: {e}')
+            logger.error(
+                f'Error sending push-notification for task_id={task.id} to URL: {url}. Error: {e}'
+            )
+            return False

tests/server/tasks/test_push_notification_sender.py

@@ -0,0 +1,133 @@
+from a2a.server.tasks.base_push_notification_sender import (
+    BasePushNotificationSender,
+)
+
+import unittest
+
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import httpx
+
+from a2a.types import (
+    PushNotificationConfig,
+    Task,
+    TaskState,
+    TaskStatus,
+)
+
+
+def create_sample_task(task_id='task123', status_state=TaskState.completed):
+    return Task(
+        id=task_id,
+        contextId='ctx456',
+        status=TaskStatus(state=status_state),
+    )
+
+
+def create_sample_push_config(
+    url='http://example.com/callback', config_id='cfg1'
+):
+    return PushNotificationConfig(id=config_id, url=url)
+
+
+class TestBasePushNotificationSender(unittest.IsolatedAsyncioTestCase):
+    def setUp(self):
+        self.mock_httpx_client = AsyncMock(spec=httpx.AsyncClient)
+        self.mock_config_store = AsyncMock()
+        self.sender = BasePushNotificationSender(
+            httpx_client=self.mock_httpx_client,
+            config_store=self.mock_config_store,
+        )
+
+    def test_constructor_stores_client_and_config_store(self):
+        self.assertEqual(self.sender._client, self.mock_httpx_client)
+        self.assertEqual(self.sender._config_store, self.mock_config_store)
+
+    async def test_send_notification_success(self):
+        task_id = 'task_send_success'
+        task_data = create_sample_task(task_id=task_id)
+        config = create_sample_push_config(url='http://notify.me/here')
+        self.mock_config_store.get_info.return_value = [config]
+
+        mock_response = AsyncMock(spec=httpx.Response)
+        mock_response.status_code = 200
+        self.mock_httpx_client.post.return_value = mock_response
+
+        await self.sender.send_notification(task_data)
+
+        self.mock_config_store.get_info.assert_awaited_once_with
+
+        # assert httpx_client post method got invoked with right parameters
+        self.mock_httpx_client.post.assert_awaited_once_with(
+            config.url,
+            json=task_data.model_dump(mode='json', exclude_none=True),
+        )
+        mock_response.raise_for_status.assert_called_once()
+
+    async def test_send_notification_no_config(self):
+        task_id = 'task_send_no_config'
+        task_data = create_sample_task(task_id=task_id)
+        self.mock_config_store.get_info.return_value = []
+
+        await self.sender.send_notification(task_data)
+
+        self.mock_config_store.get_info.assert_awaited_once_with(task_id)
+        self.mock_httpx_client.post.assert_not_called()
+
+    @patch('a2a.server.tasks.base_push_notification_sender.logger')
+    async def test_send_notification_http_status_error(
+        self, mock_logger: MagicMock
+    ):
+        task_id = 'task_send_http_err'
+        task_data = create_sample_task(task_id=task_id)
+        config = create_sample_push_config(url='http://notify.me/http_error')
+        self.mock_config_store.get_info.return_value = [config]
+
+        mock_response = MagicMock(spec=httpx.Response)
+        mock_response.status_code = 404
+        mock_response.text = 'Not Found'
+        http_error = httpx.HTTPStatusError(
+            'Not Found', request=MagicMock(), response=mock_response
+        )
+        self.mock_httpx_client.post.side_effect = http_error
+
+        await self.sender.send_notification(task_data)
+
+        self.mock_config_store.get_info.assert_awaited_once_with(task_id)
+        self.mock_httpx_client.post.assert_awaited_once_with(
+            config.url,
+            json=task_data.model_dump(mode='json', exclude_none=True),
+        )
+        mock_logger.error.assert_called_once()
+
+    async def test_send_notification_multiple_configs(self):
+        task_id = 'task_multiple_configs'
+        task_data = create_sample_task(task_id=task_id)
+        config1 = create_sample_push_config(
+            url='http://notify.me/cfg1', config_id='cfg1'
+        )
+        config2 = create_sample_push_config(
+            url='http://notify.me/cfg2', config_id='cfg2'
+        )
+        self.mock_config_store.get_info.return_value = [config1, config2]
+
+        mock_response = AsyncMock(spec=httpx.Response)
+        mock_response.status_code = 200
+        self.mock_httpx_client.post.return_value = mock_response
+
+        await self.sender.send_notification(task_data)
+
+        self.mock_config_store.get_info.assert_awaited_once_with(task_id)
+        self.assertEqual(self.mock_httpx_client.post.call_count, 2)
+
+        # Check calls for config1
+        self.mock_httpx_client.post.assert_any_call(
+            config1.url,
+            json=task_data.model_dump(mode='json', exclude_none=True),
+        )
+        # Check calls for config2
+        self.mock_httpx_client.post.assert_any_call(
+            config2.url,
+            json=task_data.model_dump(mode='json', exclude_none=True),
+        )
+        mock_response.raise_for_status.call_count = 2