"""Tests for AgentCoordinator subtask management — service layer. Uses AsyncMock for all DB operations. No real DB connections required. """ from __future__ import annotations import uuid from datetime import UTC, datetime from unittest.mock import AsyncMock, MagicMock, patch import pytest from app.plugins.builtins.automation.agent_coordinator import AgentCoordinator from app.plugins.builtins.automation.models import AgentSubtask # Override conftest DB fixtures — these tests use mocks, no real DB needed @pytest.fixture(autouse=True, scope="session") def db_setup(): """No-op override of conftest db_setup.""" yield @pytest.fixture(autouse=True) def clean_tables(db_setup): """No-op override of conftest clean_tables.""" yield @pytest.fixture(autouse=True) def clean_tables(db_setup): """No-op override of conftest clean_tables.""" yield # ─── Helpers ─── def _make_subtask( *, tenant_id: uuid.UUID | None = None, parent_agent_id: uuid.UUID | None = None, child_agent_id: uuid.UUID | None = None, task_description: str = "Do something", status: str = "pending", result: dict | None = None, ) -> AgentSubtask: """Create an AgentSubtask instance with defaults.""" return AgentSubtask( id=uuid.uuid4(), tenant_id=tenant_id or uuid.uuid4(), parent_agent_id=parent_agent_id or uuid.uuid4(), child_agent_id=child_agent_id or uuid.uuid4(), task_description=task_description, status=status, result=result or {}, created_at=datetime.now(UTC), completed_at=None, ) def _mock_session() -> AsyncMock: """Create a mock AsyncSession.""" session = AsyncMock() session.add = MagicMock() session.flush = AsyncMock() session.execute = AsyncMock() session.commit = AsyncMock() session.rollback = AsyncMock() return session # ─── Create Subtask Tests ─── class TestCreateSubtask: """Tests for AgentCoordinator.create_subtask().""" @pytest.mark.asyncio async def test_create_subtask_success(self): """create_subtask creates a subtask with pending status.""" tenant_id = uuid.uuid4() parent_id = uuid.uuid4() child_id = uuid.uuid4() db = _mock_session() def _flush_side_effect(): # After flush, the subtask should have an ID pass db.flush.side_effect = _flush_side_effect subtask = await AgentCoordinator.create_subtask( db=db, tenant_id=tenant_id, parent_agent_id=parent_id, child_agent_id=child_id, task_description="Analyze customer data", ) assert subtask.tenant_id == tenant_id assert subtask.parent_agent_id == parent_id assert subtask.child_agent_id == child_id assert subtask.task_description == "Analyze customer data" assert subtask.status == "pending" assert subtask.result == {} db.add.assert_called_once() db.flush.assert_awaited_once() @pytest.mark.asyncio async def test_create_subtask_default_status_pending(self): """create_subtask always starts with 'pending' status.""" db = _mock_session() subtask = await AgentCoordinator.create_subtask( db=db, tenant_id=uuid.uuid4(), parent_agent_id=uuid.uuid4(), child_agent_id=uuid.uuid4(), task_description="test task", ) assert subtask.status == "pending" @pytest.mark.asyncio async def test_create_subtask_default_result_empty(self): """create_subtask initializes result as empty dict.""" db = _mock_session() subtask = await AgentCoordinator.create_subtask( db=db, tenant_id=uuid.uuid4(), parent_agent_id=uuid.uuid4(), child_agent_id=uuid.uuid4(), task_description="test task", ) assert subtask.result == {} # ─── Wait For Subtask Tests ─── class TestWaitForSubtask: """Tests for AgentCoordinator.wait_for_subtask().""" @pytest.mark.asyncio async def test_wait_for_subtask_completed(self): """wait_for_subtask returns when subtask reaches 'completed' status.""" tenant_id = uuid.uuid4() db = _mock_session() subtask = _make_subtask( tenant_id=tenant_id, status="completed", result={"answer": "42"}, ) subtask.completed_at = datetime.now(UTC) mock_result = MagicMock() mock_result.scalar_one_or_none.return_value = subtask db.execute.return_value = mock_result # Patch asyncio.sleep to avoid real delays with patch("asyncio.sleep", AsyncMock()): result = await AgentCoordinator.wait_for_subtask( db=db, tenant_id=tenant_id, subtask_id=subtask.id, poll_interval=0.01, timeout=1.0, ) assert result["status"] == "completed" assert result["result"] == {"answer": "42"} assert result["subtask_id"] == str(subtask.id) @pytest.mark.asyncio async def test_wait_for_subtask_failed(self): """wait_for_subtask returns when subtask reaches 'failed' status.""" tenant_id = uuid.uuid4() subtask_id = uuid.uuid4() db = _mock_session() subtask = _make_subtask( tenant_id=tenant_id, status="failed", result={"error": "Something went wrong"}, ) subtask.completed_at = datetime.now(UTC) mock_result = MagicMock() mock_result.scalar_one_or_none.return_value = subtask db.execute.return_value = mock_result with patch("asyncio.sleep", AsyncMock()): result = await AgentCoordinator.wait_for_subtask( db=db, tenant_id=tenant_id, subtask_id=subtask_id, poll_interval=0.01, timeout=1.0, ) assert result["status"] == "failed" assert result["result"]["error"] == "Something went wrong" @pytest.mark.asyncio async def test_wait_for_subtask_cancelled(self): """wait_for_subtask returns when subtask reaches 'cancelled' status.""" tenant_id = uuid.uuid4() subtask_id = uuid.uuid4() db = _mock_session() subtask = _make_subtask( tenant_id=tenant_id, status="cancelled", result={"cancelled_by": "user"}, ) subtask.completed_at = datetime.now(UTC) mock_result = MagicMock() mock_result.scalar_one_or_none.return_value = subtask db.execute.return_value = mock_result with patch("asyncio.sleep", AsyncMock()): result = await AgentCoordinator.wait_for_subtask( db=db, tenant_id=tenant_id, subtask_id=subtask_id, poll_interval=0.01, timeout=1.0, ) assert result["status"] == "cancelled" @pytest.mark.asyncio async def test_wait_for_subtask_not_found(self): """wait_for_subtask returns error when subtask does not exist.""" db = _mock_session() mock_result = MagicMock() mock_result.scalar_one_or_none.return_value = None db.execute.return_value = mock_result with patch("asyncio.sleep", AsyncMock()): result = await AgentCoordinator.wait_for_subtask( db=db, tenant_id=uuid.uuid4(), subtask_id=uuid.uuid4(), poll_interval=0.01, timeout=1.0, ) assert result["status"] == "error" assert result["error"] == "Subtask not found" @pytest.mark.asyncio async def test_wait_for_subtask_timeout(self): """wait_for_subtask returns failed status on timeout.""" tenant_id = uuid.uuid4() subtask_id = uuid.uuid4() db = _mock_session() # Subtask stays pending forever subtask = _make_subtask(tenant_id=tenant_id, status="pending") mock_result = MagicMock() mock_result.scalar_one_or_none.return_value = subtask db.execute.return_value = mock_result with patch("asyncio.sleep", AsyncMock()): result = await AgentCoordinator.wait_for_subtask( db=db, tenant_id=tenant_id, subtask_id=subtask_id, poll_interval=0.01, timeout=0.0, # Immediate timeout ) assert result["status"] == "failed" assert "Timeout" in result["error"] # Verify the subtask was marked as failed in DB db.execute.assert_awaited() # At least the timeout update db.commit.assert_awaited() @pytest.mark.asyncio async def test_wait_for_subtask_polls_until_terminal(self): """wait_for_subtask polls multiple times before reaching terminal state.""" tenant_id = uuid.uuid4() subtask_id = uuid.uuid4() db = _mock_session() pending_subtask = _make_subtask(tenant_id=tenant_id, status="pending") completed_subtask = _make_subtask( tenant_id=tenant_id, status="completed", result={"done": True} ) completed_subtask.completed_at = datetime.now(UTC) # First poll returns pending, second returns completed mock_pending = MagicMock() mock_pending.scalar_one_or_none.return_value = pending_subtask mock_completed = MagicMock() mock_completed.scalar_one_or_none.return_value = completed_subtask db.execute.side_effect = [mock_pending, mock_completed] with patch("asyncio.sleep", AsyncMock()): result = await AgentCoordinator.wait_for_subtask( db=db, tenant_id=tenant_id, subtask_id=subtask_id, poll_interval=0.01, timeout=5.0, ) assert result["status"] == "completed" assert db.execute.await_count == 2 # Polled twice # ─── Cancel Subtask Tests ─── class TestCancelSubtask: """Tests for AgentCoordinator.cancel_subtask().""" @pytest.mark.asyncio async def test_cancel_pending_subtask(self): """cancel_subtask cancels a pending subtask.""" tenant_id = uuid.uuid4() subtask_id = uuid.uuid4() db = _mock_session() subtask = _make_subtask(tenant_id=tenant_id, status="pending") mock_result = MagicMock() mock_result.scalar_one_or_none.return_value = subtask db.execute.return_value = mock_result result = await AgentCoordinator.cancel_subtask(db, tenant_id, subtask_id) assert result is True assert subtask.status == "cancelled" assert subtask.completed_at is not None assert subtask.result["cancelled_by"] == "coordinator" db.flush.assert_awaited_once() @pytest.mark.asyncio async def test_cancel_running_subtask(self): """cancel_subtask cancels a running subtask.""" db = _mock_session() subtask = _make_subtask(status="running") mock_result = MagicMock() mock_result.scalar_one_or_none.return_value = subtask db.execute.return_value = mock_result result = await AgentCoordinator.cancel_subtask(db, uuid.uuid4(), uuid.uuid4()) assert result is True assert subtask.status == "cancelled" @pytest.mark.asyncio async def test_cancel_completed_subtask_fails(self): """cancel_subtask returns False for already completed subtask.""" db = _mock_session() subtask = _make_subtask(status="completed") mock_result = MagicMock() mock_result.scalar_one_or_none.return_value = subtask db.execute.return_value = mock_result result = await AgentCoordinator.cancel_subtask(db, uuid.uuid4(), uuid.uuid4()) assert result is False db.flush.assert_not_awaited() @pytest.mark.asyncio async def test_cancel_failed_subtask_fails(self): """cancel_subtask returns False for already failed subtask.""" db = _mock_session() subtask = _make_subtask(status="failed") mock_result = MagicMock() mock_result.scalar_one_or_none.return_value = subtask db.execute.return_value = mock_result result = await AgentCoordinator.cancel_subtask(db, uuid.uuid4(), uuid.uuid4()) assert result is False @pytest.mark.asyncio async def test_cancel_cancelled_subtask_fails(self): """cancel_subtask returns False for already cancelled subtask.""" db = _mock_session() subtask = _make_subtask(status="cancelled") mock_result = MagicMock() mock_result.scalar_one_or_none.return_value = subtask db.execute.return_value = mock_result result = await AgentCoordinator.cancel_subtask(db, uuid.uuid4(), uuid.uuid4()) assert result is False @pytest.mark.asyncio async def test_cancel_subtask_not_found(self): """cancel_subtask returns False when subtask does not exist.""" db = _mock_session() mock_result = MagicMock() mock_result.scalar_one_or_none.return_value = None db.execute.return_value = mock_result result = await AgentCoordinator.cancel_subtask(db, uuid.uuid4(), uuid.uuid4()) assert result is False # ─── Aggregate Results Tests ─── class TestAggregateResults: """Tests for AgentCoordinator.aggregate_results().""" @pytest.mark.asyncio async def test_aggregate_results_mixed_statuses(self): """aggregate_results correctly counts subtasks with mixed statuses.""" tenant_id = uuid.uuid4() db = _mock_session() sub1 = _make_subtask(status="completed", result={"output": "result1"}) sub2 = _make_subtask(status="failed", result={"error": "error1"}) sub3 = _make_subtask(status="pending") sub4 = _make_subtask(status="cancelled") mock_results = [sub1, sub2, sub3, sub4] db.execute.side_effect = [ MagicMock(scalar_one_or_none=MagicMock(return_value=st)) for st in mock_results ] result = await AgentCoordinator.aggregate_results( db=db, tenant_id=tenant_id, subtask_ids=[sub1.id, sub2.id, sub3.id, sub4.id], ) assert result["total"] == 4 assert result["completed"] == 1 assert result["failed"] == 1 assert result["pending"] == 1 assert result["cancelled"] == 1 assert len(result["results"]) == 4 @pytest.mark.asyncio async def test_aggregate_results_all_completed(self): """aggregate_results with all completed subtasks.""" db = _mock_session() sub1 = _make_subtask(status="completed", result={"a": 1}) sub2 = _make_subtask(status="completed", result={"b": 2}) db.execute.side_effect = [ MagicMock(scalar_one_or_none=MagicMock(return_value=sub1)), MagicMock(scalar_one_or_none=MagicMock(return_value=sub2)), ] result = await AgentCoordinator.aggregate_results( db=db, tenant_id=uuid.uuid4(), subtask_ids=[sub1.id, sub2.id], ) assert result["total"] == 2 assert result["completed"] == 2 assert result["failed"] == 0 @pytest.mark.asyncio async def test_aggregate_results_empty_list(self): """aggregate_results with empty subtask_ids returns zeros.""" db = _mock_session() result = await AgentCoordinator.aggregate_results( db=db, tenant_id=uuid.uuid4(), subtask_ids=[], ) assert result["total"] == 0 assert result["completed"] == 0 assert result["results"] == [] @pytest.mark.asyncio async def test_aggregate_results_with_not_found(self): """aggregate_results skips subtasks that don't exist.""" db = _mock_session() sub1 = _make_subtask(status="completed") # First subtask found, second not found db.execute.side_effect = [ MagicMock(scalar_one_or_none=MagicMock(return_value=sub1)), MagicMock(scalar_one_or_none=MagicMock(return_value=None)), ] result = await AgentCoordinator.aggregate_results( db=db, tenant_id=uuid.uuid4(), subtask_ids=[sub1.id, uuid.uuid4()], ) assert result["total"] == 1 # Only the found one assert result["completed"] == 1 # ─── List Subtasks Tests ─── class TestListSubtasks: """Tests for AgentCoordinator.list_subtasks().""" @pytest.mark.asyncio async def test_list_subtasks_no_filter(self): """list_subtasks returns all subtasks for a tenant.""" tenant_id = uuid.uuid4() db = _mock_session() sub1 = _make_subtask(tenant_id=tenant_id, status="pending") sub2 = _make_subtask(tenant_id=tenant_id, status="completed") count_result = MagicMock() count_result.scalar.return_value = 2 list_result = MagicMock() list_result.scalars.return_value.all.return_value = [sub1, sub2] db.execute.side_effect = [count_result, list_result] subtasks, total = await AgentCoordinator.list_subtasks( db=db, tenant_id=tenant_id, ) assert total == 2 assert len(subtasks) == 2 @pytest.mark.asyncio async def test_list_subtasks_filter_by_parent(self): """list_subtasks filters by parent_agent_id.""" tenant_id = uuid.uuid4() parent_id = uuid.uuid4() db = _mock_session() sub = _make_subtask(tenant_id=tenant_id, parent_agent_id=parent_id) count_result = MagicMock() count_result.scalar.return_value = 1 list_result = MagicMock() list_result.scalars.return_value.all.return_value = [sub] db.execute.side_effect = [count_result, list_result] subtasks, total = await AgentCoordinator.list_subtasks( db=db, tenant_id=tenant_id, parent_agent_id=parent_id, ) assert total == 1 assert subtasks[0].parent_agent_id == parent_id @pytest.mark.asyncio async def test_list_subtasks_filter_by_child(self): """list_subtasks filters by child_agent_id.""" tenant_id = uuid.uuid4() child_id = uuid.uuid4() db = _mock_session() sub = _make_subtask(tenant_id=tenant_id, child_agent_id=child_id) count_result = MagicMock() count_result.scalar.return_value = 1 list_result = MagicMock() list_result.scalars.return_value.all.return_value = [sub] db.execute.side_effect = [count_result, list_result] subtasks, total = await AgentCoordinator.list_subtasks( db=db, tenant_id=tenant_id, child_agent_id=child_id, ) assert total == 1 assert subtasks[0].child_agent_id == child_id @pytest.mark.asyncio async def test_list_subtasks_filter_by_status(self): """list_subtasks filters by status.""" tenant_id = uuid.uuid4() db = _mock_session() sub = _make_subtask(tenant_id=tenant_id, status="completed") count_result = MagicMock() count_result.scalar.return_value = 1 list_result = MagicMock() list_result.scalars.return_value.all.return_value = [sub] db.execute.side_effect = [count_result, list_result] subtasks, total = await AgentCoordinator.list_subtasks( db=db, tenant_id=tenant_id, status="completed", ) assert total == 1 assert subtasks[0].status == "completed" @pytest.mark.asyncio async def test_list_subtasks_pagination(self): """list_subtasks respects limit and offset.""" tenant_id = uuid.uuid4() db = _mock_session() count_result = MagicMock() count_result.scalar.return_value = 100 list_result = MagicMock() list_result.scalars.return_value.all.return_value = [] db.execute.side_effect = [count_result, list_result] subtasks, total = await AgentCoordinator.list_subtasks( db=db, tenant_id=tenant_id, limit=10, offset=20, ) assert total == 100 assert len(subtasks) == 0 @pytest.mark.asyncio async def test_list_subtasks_empty(self): """list_subtasks returns empty list when no subtasks exist.""" db = _mock_session() count_result = MagicMock() count_result.scalar.return_value = 0 list_result = MagicMock() list_result.scalars.return_value.all.return_value = [] db.execute.side_effect = [count_result, list_result] subtasks, total = await AgentCoordinator.list_subtasks( db=db, tenant_id=uuid.uuid4(), ) assert total == 0 assert len(subtasks) == 0 # ─── Model Tests ─── class TestAgentSubtaskModel: """Tests for the AgentSubtask model.""" def test_model_has_required_fields(self): """AgentSubtask has tenant_id, parent_agent_id, child_agent_id, task_description, status, result.""" subtask = AgentSubtask( tenant_id=uuid.uuid4(), parent_agent_id=uuid.uuid4(), child_agent_id=uuid.uuid4(), task_description="test task", ) assert subtask.tenant_id is not None assert subtask.parent_agent_id is not None assert subtask.child_agent_id is not None assert subtask.task_description == "test task" def test_model_default_status(self): """AgentSubtask has default='pending' for status column.""" col = AgentSubtask.__table__.c.status assert col.default.arg == "pending" def test_model_default_result(self): """AgentSubtask has default=dict for result column.""" col = AgentSubtask.__table__.c.result assert col.default.arg is dict or callable(col.default.arg) def test_model_table_name(self): """AgentSubtask uses correct table name.""" assert AgentSubtask.__tablename__ == "agent_subtasks"