backend/tests/test_agents.py

"""tests/test_agents.py — Full test suite for multi-agent system."""
import json
import pytest
from unittest.mock import MagicMock, patch, AsyncMock


# ── State tests ───────────────────────────────────────────────────────────────
class TestWorkflowState:
    def test_initial_state_created(self):
        from backend.state.graph_state import create_initial_state, TaskStatus
        state = create_initial_state("Test task")
        assert state["task"] == "Test task"
        assert state["status"] == TaskStatus.PENDING
        assert state["plan"] == []
        assert state["iteration"] == 0
        assert state["total_tokens"] == 0
        assert len(state["events"]) == 1  # task_created event

    def test_task_id_generated(self):
        from backend.state.graph_state import create_initial_state
        s1 = create_initial_state("Task A")
        s2 = create_initial_state("Task B")
        assert s1["task_id"] != s2["task_id"]

    def test_custom_task_id(self):
        from backend.state.graph_state import create_initial_state
        state = create_initial_state("Task", task_id="custom-123")
        assert state["task_id"] == "custom-123"

    def test_make_plan_step(self):
        from backend.state.graph_state import make_plan_step, StepStatus
        step = make_plan_step("s1", "Search web", "Search for X", tool="web_search", depends_on=[])
        assert step["step_id"] == "s1"
        assert step["status"] == StepStatus.PENDING
        assert step["tool"] == "web_search"
        assert step["attempts"] == 0
        assert step["result"] is None

    def test_make_agent_event(self):
        from backend.state.graph_state import make_agent_event, AgentRole
        event = make_agent_event(AgentRole.PLANNER, "plan_created", "Created 4-step plan")
        assert event["agent"] == AgentRole.PLANNER
        assert event["event_type"] == "plan_created"
        assert "timestamp" in event
        assert len(event["event_id"]) == 8


# ── Tool tests ────────────────────────────────────────────────────────────────
class TestTools:
    def test_calculate_basic(self):
        from backend.tools.registry import calculate
        result = calculate("2 + 2")
        assert result["status"] == "ok"
        assert result["data"]["result"] == 4

    def test_calculate_sqrt(self):
        from backend.tools.registry import calculate
        result = calculate("sqrt(144)")
        assert result["status"] == "ok"
        assert result["data"]["result"] == 12.0

    def test_calculate_invalid(self):
        from backend.tools.registry import calculate
        result = calculate("not_a_number()")
        assert result["status"] == "error"

    def test_write_read_file(self):
        from backend.tools.registry import write_file, read_file
        write_result = write_file("test.txt", "Hello world")
        assert write_result["status"] == "ok"
        assert write_result["data"]["filename"] == "test.txt"

        read_result = read_file("test.txt")
        assert read_result["status"] == "ok"
        assert read_result["data"]["content"] == "Hello world"

    def test_read_nonexistent_file(self):
        from backend.tools.registry import read_file
        result = read_file("does_not_exist.txt")
        assert result["status"] == "error"

    def test_run_python_basic(self):
        from backend.tools.registry import run_python
        result = run_python("print('hello')")
        assert result["status"] == "ok"
        assert "hello" in result["data"]["output"]

    def test_run_python_blocks_os(self):
        from backend.tools.registry import run_python
        result = run_python("import os")
        assert result["status"] == "error"

    def test_get_datetime(self):
        from backend.tools.registry import get_datetime
        result = get_datetime()
        assert result["status"] == "ok"
        assert "date" in result["data"]
        assert "time" in result["data"]

    @pytest.mark.asyncio
    async def test_execute_tool_dispatch(self):
        from backend.tools.registry import execute_tool
        result = await execute_tool("calculate", {"expression": "10 * 5"})
        assert result["status"] == "ok"
        assert result["data"]["result"] == 50

    @pytest.mark.asyncio
    async def test_execute_unknown_tool(self):
        from backend.tools.registry import execute_tool
        result = await execute_tool("does_not_exist", {})
        assert result["status"] == "error"
        assert "Unknown tool" in result["error"]


# ── Memory tests ──────────────────────────────────────────────────────────────
class TestShortTermMemory:
    def test_set_get(self):
        from backend.memory.memory_store import ShortTermMemory
        stm = ShortTermMemory()
        with patch("backend.memory.memory_store.get_redis") as mock_r:
            mock_redis = MagicMock()
            mock_redis.get.return_value = json.dumps({"value": 42})
            mock_r.return_value = mock_redis

            result = stm.get("task1", "key1")
            assert result == {"value": 42}

    def test_no_redis_returns_none(self):
        from backend.memory.memory_store import ShortTermMemory
        stm = ShortTermMemory()
        with patch("backend.memory.memory_store.get_redis") as mock_r:
            mock_r.return_value = None
            result = stm.get("task1", "key1")
            assert result is None


# ── Routing tests ─────────────────────────────────────────────────────────────
class TestRouting:
    def _make_state(self, status, plan=None, needs_replanning=False, iteration=0):
        from backend.state.graph_state import create_initial_state, TaskStatus
        state = create_initial_state("Test task")
        state["status"] = status
        state["plan"] = plan or []
        state["needs_replanning"] = needs_replanning
        state["iteration"] = iteration
        return state

    def test_route_executor_to_critic_when_reflecting(self):
        from backend.agents.orchestrator import route_after_executor
        from backend.state.graph_state import TaskStatus
        state = self._make_state(TaskStatus.REFLECTING)
        assert route_after_executor(state) == "critic"

    def test_route_executor_to_self_when_pending_steps(self):
        from backend.agents.orchestrator import route_after_executor
        from backend.state.graph_state import TaskStatus, StepStatus
        state = self._make_state(
            TaskStatus.EXECUTING,
            plan=[{"step_id": "s1", "status": StepStatus.PENDING}]
        )
        assert route_after_executor(state) == "executor"

    def test_route_executor_to_end_on_fatal_failure(self):
        from backend.agents.orchestrator import route_after_executor
        from backend.state.graph_state import TaskStatus
        state = self._make_state(TaskStatus.FAILED)
        assert route_after_executor(state) == "end"

    def test_route_critic_to_planner_on_replan(self):
        from backend.agents.orchestrator import route_after_critic
        from backend.state.graph_state import TaskStatus
        state = self._make_state(TaskStatus.PLANNING, needs_replanning=True)
        assert route_after_critic(state) == "planner"

    def test_route_critic_to_memory_on_approve(self):
        from backend.agents.orchestrator import route_after_critic
        from backend.state.graph_state import TaskStatus
        state = self._make_state(TaskStatus.COMPLETED)
        assert route_after_critic(state) == "memory_store"

    def test_route_planner_to_executor_with_valid_plan(self):
        from backend.agents.orchestrator import route_after_planner
        from backend.state.graph_state import TaskStatus, StepStatus
        state = self._make_state(
            TaskStatus.EXECUTING,
            plan=[{"step_id": "s1", "status": StepStatus.PENDING}]
        )
        assert route_after_planner(state) == "executor"

    def test_route_planner_to_end_with_empty_plan(self):
        from backend.agents.orchestrator import route_after_planner
        from backend.state.graph_state import TaskStatus
        state = self._make_state(TaskStatus.EXECUTING, plan=[])
        assert route_after_planner(state) == "end"


# ── API tests ─────────────────────────────────────────────────────────────────
class TestAPI:
    @pytest.fixture
    def client(self):
        from fastapi.testclient import TestClient
        from backend.api.main import app
        return TestClient(app)

    def test_health_endpoint(self, client):
        with patch("backend.api.main.get_redis", return_value=None):
            resp = client.get("/api/health")
        assert resp.status_code == 200
        data = resp.json()
        assert data["status"] == "ok"
        assert "planner" in data["agents"]

    def test_graph_endpoint(self, client):
        resp = client.get("/api/graph")
        assert resp.status_code == 200
        data = resp.json()
        assert len(data["nodes"]) == 5
        assert any(n["id"] == "planner" for n in data["nodes"])
        assert any(n["id"] == "critic" for n in data["nodes"])

    def test_task_not_found(self, client):
        resp = client.get("/api/tasks/nonexistent-id")
        assert resp.status_code == 404

    def test_task_request_validation(self, client):
        # Too short
        resp = client.post("/api/tasks", json={"task": "hi"})
        assert resp.status_code == 422

    def test_memories_endpoint(self, client):
        with patch("backend.api.main.long_term") as mock_lt:
            mock_lt.retrieve = AsyncMock(return_value=[])
            resp = client.get("/api/memories?q=test")
        assert resp.status_code == 200