Added comprehensive tests for core components

tests/test_broker.py

@@ -0,0 +1,48 @@
+import pytest
+import asyncio
+from fasta2a.broker import InMemoryBroker
+from fasta2a.schema import TaskSendParams
+
+pytestmark = pytest.mark.anyio
+
+@pytest.fixture
+def broker():
+    return InMemoryBroker()
+
+async def test_send_and_receive_task(broker: InMemoryBroker):
+    task_params: TaskSendParams = {
+        'id': 'task_123',
+        'context_id': 'ctx_abc',
+        'message': {
+            'role': 'user',
+            'parts': [{'kind': 'text', 'text': 'test message'}],
+            'kind': 'message',
+            'message_id': 'msg_xyz'
+        }
+    }
+
+    result_queue = asyncio.Queue()
+
+    async def consumer():
+        async for op in broker.receive_task_operations():
+            await result_queue.put(op)
+            break
+
+    async with broker:
+        consumer_task = asyncio.create_task(consumer())
+        await asyncio.sleep(0.01)
+
+        await broker.run_task(task_params)
+
+        try:
+            received_op = await asyncio.wait_for(result_queue.get(), timeout=1)
+        finally:
+            consumer_task.cancel()
+            try:
+                await consumer_task
+            except asyncio.CancelledError:
+                pass
+
+    assert received_op is not None
+    assert received_op['operation'] == 'run'
+    assert received_op['params']['id'] == 'task_123'

tests/test_storage.py

@@ -0,0 +1,55 @@
+import pytest
+from fasta2a.storage import InMemoryStorage, ContextT
+from fasta2a.schema import Message, TextPart
+
+pytestmark = pytest.mark.anyio
+
+@pytest.fixture
+def storage() -> InMemoryStorage[ContextT]:
+    return InMemoryStorage()
+
+@pytest.mark.asyncio
+async def test_submit_and_load_task(storage: InMemoryStorage):
+    message = Message(
+        role='user',
+        parts=[TextPart(text='hello', kind='text')],
+        kind='message',
+        message_id='msg1'
+    )
+
+    task = await storage.submit_task(context_id="ctx1", message=message)
+    assert task is not None
+    assert task['id'] is not None
+    assert task['status']['state'] == 'submitted'
+    assert task['context_id'] == 'ctx1'
+
+    loaded_task = await storage.load_task(task['id'])
+    assert loaded_task is not None
+    assert loaded_task['id'] == task['id']
+
+@pytest.mark.asyncio
+async def test_update_task(storage: InMemoryStorage):
+    message = Message(
+        role='user',
+        parts=[TextPart(text='hello', kind='text')],
+        kind='message',
+        message_id='msg2'
+    )
+    task = await storage.submit_task(context_id="ctx2", message=message)
+
+    await storage.update_task(task['id'], state='completed')
+
+    updated_task = await storage.load_task(task['id'])
+    assert updated_task is not None
+    assert updated_task['status']['state'] == 'completed'
+
+@pytest.mark.asyncio
+async def test_context_storage(storage: InMemoryStorage):
+    context = await storage.load_context("ctx3")
+    assert context is None
+
+    await storage.update_context("ctx3", ["message1", "message2"])
+
+    loaded_context = await storage.load_context("ctx3")
+    assert loaded_context is not None
+    assert loaded_context == ["message1", "message2"]

tests/test_task_manager.py

@@ -0,0 +1,105 @@
+# tests/test_task_manager.py
+
+import pytest
+import asyncio
+from typing import Any
+
+from fasta2a.task_manager import TaskManager
+from fasta2a.broker import InMemoryBroker
+from fasta2a.storage import InMemoryStorage, ContextT
+from fasta2a.worker import Worker
+from fasta2a.schema import (
+    Artifact,
+    Message,
+    TextPart,
+    TaskSendParams,
+    TaskIdParams,
+    SendMessageRequest,
+    GetTaskRequest
+)
+
+pytestmark = pytest.mark.anyio
+
+class MockWorker(Worker[ContextT]):
+    def __init__(self, storage, broker):
+        super().__init__(storage=storage, broker=broker)
+        self.executed_tasks = []
+        self.canceled_tasks = []
+        self.is_running = asyncio.Event()
+
+    async def run_task(self, params: TaskSendParams) -> None:
+        self.executed_tasks.append(params['id'])
+        await self.storage.update_task(params['id'], state='working')
+        await asyncio.sleep(0.1)
+        response_message = Message(
+            role='agent',
+            parts=[TextPart(text=f"Completed task {params['id']}", kind='text')],
+            kind='message',
+            message_id='response_msg'
+        )
+        await self.storage.update_task(
+            params['id'],
+            state='completed',
+            new_messages=[response_message]
+        )
+
+    async def cancel_task(self, params: TaskIdParams) -> None:
+        self.canceled_tasks.append(params['id'])
+        await self.storage.update_task(params['id'], state='canceled')
+
+    def build_message_history(self, history: list[Message]) -> list[Any]:
+        return []
+
+    def build_artifacts(self, result: Any) -> list[Artifact]:
+        return []
+
+    async def _loop(self) -> None:
+        self.is_running.set()
+        await super()._loop()
+
+@pytest.fixture
+def storage() -> InMemoryStorage[ContextT]:
+    return InMemoryStorage()
+
+@pytest.fixture
+def broker():
+    return InMemoryBroker()
+
+@pytest.fixture
+def worker(storage, broker):
+    return MockWorker(storage=storage, broker=broker)
+
+@pytest.fixture
+def task_manager(storage, broker):
+    return TaskManager(storage=storage, broker=broker)
+
+async def test_task_lifecycle(task_manager: TaskManager, worker: MockWorker, storage: InMemoryStorage):
+    async with task_manager, worker.run():
+        await worker.is_running.wait()
+        send_request = SendMessageRequest(
+            jsonrpc='2.0',
+            id='req1',
+            method='message/send',
+            params={
+                'message': {
+                    'role': 'user',
+                    'parts': [{'kind': 'text', 'text': 'Please run a test task.'}],
+                    'kind': 'message',
+                    'message_id': 'user_msg1'
+                }
+            }
+        )
+        response = await task_manager.send_message(send_request)
+        task_id = response['result']['id']
+        await asyncio.sleep(0.2)
+        assert task_id in worker.executed_tasks
+        get_request = GetTaskRequest(
+            jsonrpc='2.0',
+            id='req2',
+            method='tasks/get',
+            params={'id': task_id}
+        )
+        get_response = await task_manager.get_task(get_request)
+        assert get_response['result']['status']['state'] == 'completed'
+        assert len(get_response['result']['history']) == 2
+        assert get_response['result']['history'][-1]['role'] == 'agent'