Verify that notification tokens were properly passed around

Author: lkawka

tests/e2e/push_notifications/notifications_app.py

@@ -3,34 +3,49 @@
 from typing import Annotated
 
 from fastapi import FastAPI, HTTPException, Path, Request
+from pydantic import BaseModel, ValidationError
+
+from a2a.types import Task
+
+
+class Notification(BaseModel):
+    """Encapsulates default push notification data."""
+
+    task: Task
+    token: str
 
 
 def create_notifications_app() -> FastAPI:
     """Creates a simple push notification ingesting HTTP+REST application."""
     app = FastAPI()
     store_lock = asyncio.Lock()
-    store: dict[str, list] = {}
+    store: dict[str, list[Notification]] = {}
 
     @app.post('/notifications')
     async def add_notification(request: Request):
         """Endpoint for injesting notifications from agents. It receives a JSON
         payload and stores it in-memory.
         """
-        if not request.headers.get('x-a2a-notification-token'):
+        token = request.headers.get('x-a2a-notification-token')
+        if not token:
             raise HTTPException(
                 status_code=400,
                 detail='Missing "x-a2a-notification-token" header.',
             )
-        payload = await request.json()
-        task_id = payload.get('id')
-        if not task_id:
-            raise HTTPException(
-                status_code=400, detail='Missing "id" in notification payload.'
-            )
+        try:
+            task = Task.model_validate(await request.json())
+        except ValidationError as e:
+            raise HTTPException(status_code=400, detail=str(e))
+
         async with store_lock:
-            if task_id not in store:
-                store[task_id] = []
-            store[task_id].append(payload)
+            if task.id not in store:
+                store[task.id] = []
+            store[task.id].append(
+                Notification(
+                    task=task,
+                    token=token,
+                )
+            )
         return {
             'status': 'received',
         }

tests/e2e/push_notifications/test_default_push_notification_support.py

@@ -9,7 +9,7 @@
 import pytest_asyncio
 
 from agent_app import create_agent_app
-from notifications_app import create_notifications_app
+from notifications_app import Notification, create_notifications_app
 from utils import (
     create_app_process,
     find_free_port,
@@ -115,14 +115,15 @@ async def test_notification_triggering_with_in_message_config_e2e(
     Tests push notification triggering for in-message push notification config.
     """
     # Create an A2A client with a push notification config.
+    token = uuid.uuid4().hex
     a2a_client = ClientFactory(
         ClientConfig(
             supported_transports=[TransportProtocol.http_json],
             push_notification_configs=[
                 PushNotificationConfig(
                     id='in-message-config',
                     url=f'{notifications_server}/notifications',
-                    token=uuid.uuid4().hex,
+                    token=token,
                 )
             ],
         )
@@ -150,7 +151,9 @@ async def test_notification_triggering_with_in_message_config_e2e(
         f'{notifications_server}/tasks/{task.id}/notifications',
         n=1,
     )
-    assert notifications[0].status.state == 'completed'
+    assert notifications[0].token == token
+    assert notifications[0].task.id == task.id
+    assert notifications[0].task.status.state == 'completed'
 
 
 @pytest.mark.asyncio
@@ -192,13 +195,14 @@ async def test_notification_triggering_after_config_change_e2e(
     assert len(response.json().get('notifications', [])) == 0
 
     # Set the push notification config.
+    token = uuid.uuid4().hex
     await a2a_client.set_task_callback(
         TaskPushNotificationConfig(
             task_id=task.id,
             push_notification_config=PushNotificationConfig(
                 id='after-config-change',
                 url=f'{notifications_server}/notifications',
-                token=uuid.uuid4().hex,
+                token=token,
             ),
         )
     )
@@ -223,15 +227,17 @@ async def test_notification_triggering_after_config_change_e2e(
         f'{notifications_server}/tasks/{task.id}/notifications',
         n=1,
     )
-    assert notifications[0].status.state == 'completed'
+    assert notifications[0].task.id == task.id
+    assert notifications[0].task.status.state == 'completed'
+    assert notifications[0].token == token
 
 
 async def wait_for_n_notifications(
     http_client: httpx.AsyncClient,
     url: str,
     n: int,
     timeout: int = 3,
-) -> list[Task]:
+) -> list[Notification]:
     """
     Queries the notification URL until the desired number of notifications
     is received or the timeout is reached.
@@ -243,7 +249,7 @@ async def wait_for_n_notifications(
         assert response.status_code == 200
         notifications = response.json()['notifications']
         if len(notifications) == n:
-            return [Task.model_validate(n) for n in notifications]
+            return [Notification.model_validate(n) for n in notifications]
         if time.time() - start_time > timeout:
             raise TimeoutError(
                 f'Notification retrieval timed out. Got {len(notifications)} notification(s), want {n}. Retrieved notifications: {notifications}.'