tests/test_workflows.py

import pytest
from agents.base_agent import AgentStatus, AgentResult
from workflows.sequential_workflow import SequentialWorkflow
from workflows.parallel_workflow import ParallelWorkflow
from workflows.hybrid_workflow import HybridWorkflow
from typing import Dict, Any


class DummyAgent:
    def __init__(self, name, wave=1, critical=True, fail=False):
        self.name = name
        self.wave = wave
        self.critical = critical
        self._fail = fail

    async def run(self, ctx) -> AgentResult:
        if self._fail:
            return AgentResult(
                agent_name=self.name, status=AgentStatus.FAILED, error="test error"
            )
        return AgentResult(
            agent_name=self.name, status=AgentStatus.COMPLETED, data={"done": True}
        )


@pytest.mark.asyncio
async def test_sequential_runs_in_order(agent_ctx):
    agents = [DummyAgent("a", wave=1), DummyAgent("b", wave=2)]
    results = await SequentialWorkflow().execute(agents, agent_ctx)
    assert len(results) == 2
    assert all(r.status == AgentStatus.COMPLETED for r in results)


@pytest.mark.asyncio
async def test_parallel_workflow_runs_all(agent_ctx):
    agents = [DummyAgent("x", wave=2), DummyAgent("y", wave=2), DummyAgent("z", wave=2)]
    results = await ParallelWorkflow().execute(agents, agent_ctx)
    assert len(results) == 3
    assert all(r.status == AgentStatus.COMPLETED for r in results)


@pytest.mark.asyncio
async def test_hybrid_stops_on_critical_failure(agent_ctx):
    agents = [
        DummyAgent("parser", wave=1, critical=True, fail=True),
        DummyAgent("claim", wave=2, critical=True),
    ]
    results = await HybridWorkflow().execute(agents, agent_ctx)
    # Wave 2 should not have run
    names = [r.agent_name for r in results]
    assert "claim" not in names


@pytest.mark.asyncio
async def test_hybrid_continues_after_non_critical_failure(agent_ctx):
    agents = [
        DummyAgent("parser", wave=1, critical=True),
        DummyAgent("keyword", wave=2, critical=False, fail=True),
        DummyAgent("claim", wave=2, critical=True),
        DummyAgent("scoring", wave=4, critical=True),
    ]
    results = await HybridWorkflow().execute(agents, agent_ctx)
    names = [r.agent_name for r in results]
    assert "scoring" in names  # pipeline continued