tests/test_workflow_integration.py

import ast
import asyncio
import sys
import types
import unittest
from enum import IntEnum
from pathlib import Path

from pydantic import ValidationError

BACKEND_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(BACKEND_ROOT))


def _install_auth_stub() -> None:
    auth_guard = types.ModuleType("api.auth_guard")

    class AuthRole(IntEnum):
        MACHINE = 1

    auth_guard.AuthRole = AuthRole
    auth_guard.require_role = lambda _role: (lambda: None)
    sys.modules["api.auth_guard"] = auth_guard


_install_auth_stub()

from agents.workflow_engine import Workflow, WorkflowExecutor, WorkflowStep
from api import workflows


class FakeKernel:
    def __init__(self, resolutions):
        self.resolutions = resolutions
        self.calls = []

    async def resolve_capability(self, tool_name):
        self.calls.append(tool_name)
        return self.resolutions.get(tool_name, {"status": "error"})


class FakeExecutor:
    def __init__(self, results=None):
        self.results = results or {}
        self.calls = []

    async def run_tool(self, tool_name, inputs, timeout=30.0, worker_hint=None):
        self.calls.append({
            "tool_name": tool_name,
            "inputs": inputs,
            "timeout": timeout,
            "worker_hint": worker_hint,
        })
        result = self.results.get(tool_name, {"success": True, "output": tool_name})
        if isinstance(result, Exception):
            raise result
        return result


class WorkflowExecutorTests(unittest.IsolatedAsyncioTestCase):
    async def test_uses_resolved_worker_hint_and_records_completed_step(self):
        kernel = FakeKernel({"search": {"status": "resolved", "worker": {"id": "worker-eu"}}})
        executor = FakeExecutor()
        engine = WorkflowExecutor(kernel=kernel, executor=executor)
        workflow = Workflow(name="ricerca", steps=[WorkflowStep(tool_name="search", args={"query": "AI"})])

        result = await engine.execute_workflow(workflow)

        self.assertEqual(result.status, "completed")
        self.assertEqual(result.steps[0].status, "completed")
        self.assertIsNotNone(result.steps[0].started_at)
        self.assertIsNotNone(result.steps[0].finished_at)
        self.assertEqual(executor.calls[0]["worker_hint"], "worker-eu")

    async def test_preserves_local_fallback_when_no_worker_is_resolved(self):
        kernel = FakeKernel({})
        executor = FakeExecutor()
        engine = WorkflowExecutor(kernel=kernel, executor=executor)
        workflow = Workflow(name="fallback", steps=[WorkflowStep(tool_name="local", args={"value": 1})])

        result = await engine.execute_workflow(workflow)

        self.assertEqual(result.status, "completed")
        self.assertIsNone(executor.calls[0]["worker_hint"])

    async def test_stops_after_unsuccessful_tool_result(self):
        kernel = FakeKernel({})
        executor = FakeExecutor({"first": {"success": False, "error": "denied"}})
        engine = WorkflowExecutor(kernel=kernel, executor=executor)
        workflow = Workflow(
            name="errore",
            steps=[
                WorkflowStep(tool_name="first", args={}),
                WorkflowStep(tool_name="second", args={}),
            ],
        )

        result = await engine.execute_workflow(workflow)

        self.assertEqual(result.status, "failed")
        self.assertEqual(result.steps[0].status, "failed")
        self.assertEqual(result.steps[0].error, "denied")
        self.assertEqual(result.steps[1].status, "pending")
        self.assertEqual([call["tool_name"] for call in executor.calls], ["first"])


class WorkflowApiTests(unittest.IsolatedAsyncioTestCase):
    def setUp(self):
        self.previous_executor = workflows._workflow_executor
        self.previous_tasks = workflows._workflow_tasks.copy()
        workflows._workflow_executor = None
        workflows._workflow_tasks.clear()

    async def asyncTearDown(self):
        for task in list(workflows._workflow_tasks.values()):
            task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass
        workflows._workflow_tasks.clear()
        workflows._workflow_executor = self.previous_executor
        workflows._workflow_tasks.update(self.previous_tasks)

    def test_router_exposes_start_and_status_paths(self):
        routes = {
            (route.path, method)
            for route in workflows.router.routes
            for method in (getattr(route, "methods", set()) or set())
        }

        self.assertEqual(workflows.router.prefix, "/api/workflows")
        self.assertIn(("/api/workflows", "POST"), routes)
        self.assertIn(("/api/workflows/{workflow_id}", "GET"), routes)

    def test_main_router_map_mounts_workflows(self):
        tree = ast.parse((BACKEND_ROOT / "main.py").read_text(encoding="utf-8"))
        router_map = next(
            node.value
            for node in tree.body
            if isinstance(node, ast.Assign)
            and any(isinstance(target, ast.Name) and target.id == "_ROUTER_MAP" for target in node.targets)
        )
        routes = {
            key.value: value.value
            for key, value in zip(router_map.keys, router_map.values)
            if isinstance(key, ast.Constant) and isinstance(value, ast.Constant)
        }

        self.assertEqual(routes.get("workflows"), "workflows")

    def test_start_contract_rejects_empty_steps(self):
        with self.assertRaises(ValidationError):
            workflows.WorkflowStartIn(name="vuoto", steps=[])

    async def test_start_and_get_workflow_use_background_execution(self):
        engine = WorkflowExecutor(kernel=FakeKernel({}), executor=FakeExecutor())
        workflows._workflow_executor = engine
        body = workflows.WorkflowStartIn(
            name="API workflow",
            steps=[WorkflowStep(tool_name="local", args={"n": 1})],
        )

        started = await workflows.start_workflow(body)
        task = workflows._workflow_tasks[started.workflow_id]
        await task
        fetched = await workflows.get_workflow(started.workflow_id)

        self.assertEqual(fetched.workflow_id, started.workflow_id)
        self.assertEqual(fetched.status, "completed")
        self.assertEqual(fetched.steps[0].status, "completed")


if __name__ == "__main__":
    unittest.main()