Spaces:
Running
Running
| import asyncio | |
| import os | |
| import time | |
| from unittest.mock import AsyncMock, MagicMock | |
| import pytest | |
| from helpers.blocking_a2a import BlockingTaskManager | |
| pytestmark = pytest.mark.asyncio | |
| def _make_request(task_id="task-123", blocking=False): | |
| return { | |
| "jsonrpc": "2.0", | |
| "id": "req-1", | |
| "method": "message/send", | |
| "params": { | |
| "message": { | |
| "kind": "message", | |
| "role": "user", | |
| "message_id": task_id, | |
| "parts": [{"kind": "text", "text": "hello"}], | |
| }, | |
| "configuration": {"blocking": blocking}, | |
| }, | |
| } | |
| def _make_super_response(task_id="task-123", state="submitted"): | |
| return { | |
| "jsonrpc": "2.0", | |
| "id": "req-1", | |
| "result": { | |
| "id": task_id, | |
| "status": {"state": state}, | |
| }, | |
| } | |
| async def test_non_blocking_returns_immediately(monkeypatch): | |
| broker = MagicMock() | |
| storage = MagicMock() | |
| tm = BlockingTaskManager(broker=broker, storage=storage) | |
| req = _make_request(blocking=False) | |
| super_resp = _make_super_response("task-123", "submitted") | |
| monkeypatch.setattr( | |
| "fasta2a.task_manager.TaskManager.send_message", | |
| AsyncMock(return_value=super_resp), | |
| ) | |
| resp = await tm.send_message(req) | |
| assert resp == super_resp | |
| storage.load_task.assert_not_called() | |
| async def test_blocking_immediate_completion(monkeypatch): | |
| broker = MagicMock() | |
| storage = MagicMock() | |
| tm = BlockingTaskManager(broker=broker, storage=storage) | |
| req = _make_request(blocking=True) | |
| super_resp = _make_super_response("task-101", "submitted") | |
| completed_task = { | |
| "id": "task-101", | |
| "status": {"state": "completed"}, | |
| "history": [], | |
| } | |
| storage.load_task = AsyncMock(return_value=completed_task) | |
| monkeypatch.setattr( | |
| "fasta2a.task_manager.TaskManager.send_message", | |
| AsyncMock(return_value=super_resp), | |
| ) | |
| resp = await tm.send_message(req) | |
| assert resp["jsonrpc"] == "2.0" | |
| assert resp["id"] == "req-1" | |
| assert resp["result"]["status"]["state"] == "completed" | |
| async def test_blocking_delayed_completion(monkeypatch): | |
| broker = MagicMock() | |
| storage = MagicMock() | |
| tm = BlockingTaskManager(broker=broker, storage=storage) | |
| req = _make_request(blocking=True) | |
| super_resp = _make_super_response("task-202", "submitted") | |
| running_task = { | |
| "id": "task-202", | |
| "status": {"state": "running"}, | |
| "history": [], | |
| } | |
| completed_task = { | |
| "id": "task-202", | |
| "status": {"state": "completed"}, | |
| "history": [], | |
| } | |
| states = iter([running_task, running_task, completed_task]) | |
| async def mock_load_task(task_id): | |
| return next(states) | |
| storage.load_task = mock_load_task | |
| monkeypatch.setattr( | |
| "fasta2a.task_manager.TaskManager.send_message", | |
| AsyncMock(return_value=super_resp), | |
| ) | |
| start_time = time.monotonic() | |
| resp = await tm.send_message(req) | |
| elapsed = time.monotonic() - start_time | |
| assert resp["result"]["status"]["state"] == "completed" | |
| assert elapsed >= 0.9 # 2 polls * 0.5s sleep | |
| async def test_blocking_failure_state(monkeypatch): | |
| broker = MagicMock() | |
| storage = MagicMock() | |
| tm = BlockingTaskManager(broker=broker, storage=storage) | |
| req = _make_request(blocking=True) | |
| super_resp = _make_super_response("task-303", "submitted") | |
| failed_task = { | |
| "id": "task-303", | |
| "status": {"state": "failed"}, | |
| "history": [], | |
| } | |
| storage.load_task = AsyncMock(return_value=failed_task) | |
| monkeypatch.setattr( | |
| "fasta2a.task_manager.TaskManager.send_message", | |
| AsyncMock(return_value=super_resp), | |
| ) | |
| resp = await tm.send_message(req) | |
| assert resp["result"]["status"]["state"] == "failed" | |
| async def test_blocking_timeout_returns_current_task(monkeypatch): | |
| broker = MagicMock() | |
| storage = MagicMock() | |
| tm = BlockingTaskManager(broker=broker, storage=storage) | |
| req = _make_request(blocking=True) | |
| super_resp = _make_super_response("task-404", "submitted") | |
| working_task = { | |
| "id": "task-404", | |
| "status": {"state": "running"}, | |
| "history": [], | |
| } | |
| storage.load_task = AsyncMock(return_value=working_task) | |
| monkeypatch.setenv("A2A_BLOCKING_TIMEOUT", "1") | |
| monkeypatch.setattr( | |
| "fasta2a.task_manager.TaskManager.send_message", | |
| AsyncMock(return_value=super_resp), | |
| ) | |
| start_time = time.monotonic() | |
| resp = await tm.send_message(req) | |
| elapsed = time.monotonic() - start_time | |
| assert resp["result"]["status"]["state"] == "running" | |
| assert 0.9 <= elapsed <= 2.5 | |
| async def test_concurrent_blocking_requests(monkeypatch): | |
| broker = MagicMock() | |
| storage = MagicMock() | |
| tm = BlockingTaskManager(broker=broker, storage=storage) | |
| req_sync = _make_request("task-sync", blocking=True) | |
| req_async = _make_request("task-async", blocking=False) | |
| super_resp_sync = _make_super_response("task-sync", "submitted") | |
| super_resp_async = _make_super_response("task-async", "submitted") | |
| async def mock_send_message(self, request): | |
| task_id = request["params"]["message"]["message_id"] | |
| if request["params"].get("configuration", {}).get("blocking"): | |
| return _make_super_response(task_id, "submitted") | |
| return _make_super_response(task_id, "submitted") | |
| async def mock_load_task(task_id): | |
| if task_id == "task-sync": | |
| return {"id": "task-sync", "status": {"state": "completed"}} | |
| return {"id": "task-async", "status": {"state": "submitted"}} | |
| storage.load_task = mock_load_task | |
| monkeypatch.setattr( | |
| "fasta2a.task_manager.TaskManager.send_message", | |
| mock_send_message, | |
| ) | |
| res_sync, res_async = await asyncio.gather( | |
| tm.send_message(req_sync), | |
| tm.send_message(req_async), | |
| ) | |
| assert res_sync["result"]["status"]["state"] == "completed" | |
| assert res_async["result"]["status"]["state"] == "submitted" | |