ForcePilot/backend/test/test_subagent.py

"""SubAgent 单元测试"""

from __future__ import annotations

from contextlib import asynccontextmanager
from unittest.mock import AsyncMock, MagicMock

import pytest
from sqlalchemy.exc import IntegrityError

from yuxi.storage.postgres.models_business import SubAgent
from yuxi.utils.datetime_utils import utc_now_naive


# =============================================================================
# Router Tests
# =============================================================================

from fastapi import FastAPI, HTTPException
from fastapi.testclient import TestClient

from server.routers.subagent_router import subagents_router
from server.utils.auth_middleware import get_admin_user, get_db
from yuxi.storage.postgres.models_business import User


def _build_app() -> FastAPI:
    app = FastAPI()
    app.include_router(subagents_router, prefix="/api")

    async def fake_db():
        return None

    async def fake_admin_user():
        return User(
            username="admin",
            user_id="admin",
            password_hash="x",
            role="admin",
        )

    app.dependency_overrides[get_db] = fake_db
    app.dependency_overrides[get_admin_user] = fake_admin_user
    return app


def test_list_subagents_returns_data(monkeypatch):
    async def fake_get_all_subagents(_db):
        return [
            {
                "name": "research-agent",
                "description": "Test research agent",
                "system_prompt": "You are a researcher",
                "tools": ["tavily_search"],
                "model": None,
                "is_builtin": True,
                "created_by": "system",
                "updated_by": "system",
                "created_at": "2024-01-01T00:00:00",
                "updated_at": "2024-01-01T00:00:00",
            }
        ]

    monkeypatch.setattr("server.routers.subagent_router.service.get_all_subagents", fake_get_all_subagents)

    app = _build_app()
    client = TestClient(app)
    resp = client.get("/api/system/subagents")
    assert resp.status_code == 200, resp.text
    payload = resp.json()
    assert payload["success"] is True
    assert payload["data"][0]["name"] == "research-agent"
    assert payload["data"][0]["is_builtin"] is True


def test_get_single_subagent(monkeypatch):
    async def fake_get_subagent(name, db=None):
        if name == "research-agent":
            return {
                "name": "research-agent",
                "description": "Test research agent",
                "system_prompt": "You are a researcher",
                "tools": ["tavily_search"],
                "model": None,
                "is_builtin": True,
                "created_by": "system",
                "updated_by": "system",
                "created_at": "2024-01-01T00:00:00",
                "updated_at": "2024-01-01T00:00:00",
            }
        return None

    monkeypatch.setattr("server.routers.subagent_router.service.get_subagent", fake_get_subagent)

    app = _build_app()
    client = TestClient(app)
    resp = client.get("/api/system/subagents/research-agent")
    assert resp.status_code == 200, resp.text
    payload = resp.json()
    assert payload["success"] is True
    assert payload["data"]["name"] == "research-agent"


def test_get_single_subagent_not_found(monkeypatch):
    async def fake_get_subagent(name, db=None):
        return None

    monkeypatch.setattr("server.routers.subagent_router.service.get_subagent", fake_get_subagent)

    app = _build_app()
    client = TestClient(app)
    resp = client.get("/api/system/subagents/nonexistent")
    assert resp.status_code == 404, resp.text


def test_create_subagent(monkeypatch):
    captured = {}

    async def fake_create_subagent(data, created_by, db=None):
        captured["data"] = data
        captured["created_by"] = created_by
        return {
            "name": data["name"],
            "description": data["description"],
            "system_prompt": data["system_prompt"],
            "tools": data.get("tools", []),
            "model": data.get("model"),
            "is_builtin": False,
            "created_by": created_by,
            "updated_by": created_by,
            "created_at": "2024-01-01T00:00:00",
            "updated_at": "2024-01-01T00:00:00",
        }

    monkeypatch.setattr("server.routers.subagent_router.service.create_subagent", fake_create_subagent)

    app = _build_app()
    client = TestClient(app)
    resp = client.post(
        "/api/system/subagents",
        json={
            "name": "my-agent",
            "description": "My custom agent",
            "system_prompt": "You are a helpful assistant",
            "tools": ["tool_a", "tool_b"],
            "model": None,
        },
    )
    assert resp.status_code == 200, resp.text
    payload = resp.json()
    assert payload["success"] is True
    assert captured["data"]["name"] == "my-agent"
    assert captured["created_by"] == "admin"


def test_create_subagent_duplicate_returns_409(monkeypatch):
    async def fake_create_subagent(data, created_by, db=None):
        raise IntegrityError(
            "duplicate",
            {},
            Exception('duplicate key value violates unique constraint "subagents_pkey"'),
        )

    monkeypatch.setattr("server.routers.subagent_router.service.create_subagent", fake_create_subagent)

    app = _build_app()
    client = TestClient(app)
    resp = client.post(
        "/api/system/subagents",
        json={
            "name": "my-agent",
            "description": "My custom agent",
            "system_prompt": "You are a helpful assistant",
            "tools": [],
            "model": None,
        },
    )
    assert resp.status_code == 409, resp.text


def test_update_subagent(monkeypatch):
    captured = {}

    async def fake_update_subagent(name, data, updated_by, db=None):
        captured["name"] = name
        captured["data"] = data
        captured["updated_by"] = updated_by
        return {
            "name": name,
            "description": data.get("description", "Updated description"),
            "system_prompt": data.get("system_prompt", "Updated prompt"),
            "tools": data.get("tools", []),
            "model": data.get("model"),
            "is_builtin": False,
            "created_by": "admin",
            "updated_by": updated_by,
            "created_at": "2024-01-01T00:00:00",
            "updated_at": "2024-01-01T00:00:00",
        }

    monkeypatch.setattr("server.routers.subagent_router.service.update_subagent", fake_update_subagent)

    app = _build_app()
    client = TestClient(app)
    resp = client.put(
        "/api/system/subagents/my-agent",
        json={
            "description": "Updated description",
            "system_prompt": "Updated prompt",
        },
    )
    assert resp.status_code == 200, resp.text
    payload = resp.json()
    assert payload["success"] is True
    assert captured["name"] == "my-agent"
    assert captured["updated_by"] == "admin"


def test_update_builtin_subagent_fails(monkeypatch):
    async def fake_update_subagent(name, data, updated_by, db=None):
        raise ValueError("内置 SubAgent 不可编辑")

    monkeypatch.setattr("server.routers.subagent_router.service.update_subagent", fake_update_subagent)

    app = _build_app()
    client = TestClient(app)
    resp = client.put(
        "/api/system/subagents/research-agent",
        json={"description": "Try to update builtin"},
    )
    assert resp.status_code == 400, resp.text


def test_delete_subagent(monkeypatch):
    deleted_name = {"name": None}

    async def fake_delete_subagent(name, db=None):
        deleted_name["name"] = name
        return True

    monkeypatch.setattr("server.routers.subagent_router.service.delete_subagent", fake_delete_subagent)

    app = _build_app()
    client = TestClient(app)
    resp = client.delete("/api/system/subagents/my-agent")
    assert resp.status_code == 200, resp.text
    assert deleted_name["name"] == "my-agent"


def test_delete_builtin_subagent_fails(monkeypatch):
    async def fake_delete_subagent(name, db=None):
        raise ValueError("内置 SubAgent 不可删除")

    monkeypatch.setattr("server.routers.subagent_router.service.delete_subagent", fake_delete_subagent)

    app = _build_app()
    client = TestClient(app)
    resp = client.delete("/api/system/subagents/research-agent")
    assert resp.status_code == 400, resp.text


# =============================================================================
# Repository Tests
# =============================================================================

class TestSubAgentRepository:
    @pytest.mark.asyncio
    async def test_list_all(self):
        from yuxi.repositories.subagent_repository import SubAgentRepository

        mock_db = AsyncMock()
        mock_result = MagicMock()
        mock_result.scalars.return_value.all.return_value = [
            SubAgent(
                name="test-agent",
                description="Test agent",
                system_prompt="You are a test",
                tools=["tool_a"],
                model=None,
                is_builtin=False,
                created_by="admin",
                updated_by="admin",
                created_at=utc_now_naive(),
                updated_at=utc_now_naive(),
            )
        ]
        mock_db.execute.return_value = mock_result

        repo = SubAgentRepository(mock_db)
        result = await repo.list_all()

        assert len(result) == 1
        assert result[0].name == "test-agent"
        mock_db.execute.assert_called_once()

    @pytest.mark.asyncio
    async def test_get_by_name_found(self):
        from yuxi.repositories.subagent_repository import SubAgentRepository

        mock_db = AsyncMock()
        mock_result = MagicMock()
        mock_result.scalar_one_or_none.return_value = SubAgent(
            name="test-agent",
            description="Test agent",
            system_prompt="You are a test",
            tools=[],
            model=None,
            is_builtin=False,
            created_by="admin",
            updated_by="admin",
            created_at=utc_now_naive(),
            updated_at=utc_now_naive(),
        )
        mock_db.execute.return_value = mock_result

        repo = SubAgentRepository(mock_db)
        result = await repo.get_by_name("test-agent")

        assert result is not None
        assert result.name == "test-agent"

    @pytest.mark.asyncio
    async def test_get_by_name_not_found(self):
        from yuxi.repositories.subagent_repository import SubAgentRepository

        mock_db = AsyncMock()
        mock_result = MagicMock()
        mock_result.scalar_one_or_none.return_value = None
        mock_db.execute.return_value = mock_result

        repo = SubAgentRepository(mock_db)
        result = await repo.get_by_name("nonexistent")

        assert result is None

    @pytest.mark.asyncio
    async def test_update_can_clear_model_when_provided(self):
        from yuxi.repositories.subagent_repository import SubAgentRepository

        mock_db = AsyncMock()
        repo = SubAgentRepository(mock_db)
        item = SubAgent(
            name="test-agent",
            description="Test agent",
            system_prompt="You are a test",
            tools=[],
            model="gpt-4",
            is_builtin=False,
            created_by="admin",
            updated_by="admin",
            created_at=utc_now_naive(),
            updated_at=utc_now_naive(),
        )

        await repo.update(
            item,
            description=None,
            system_prompt=None,
            tools=None,
            model=None,
            model_provided=True,
            updated_by="admin",
        )

        assert item.model is None


# =============================================================================
# Service Tests
# =============================================================================

class TestSubAgentService:
    @pytest.mark.asyncio
    async def test_init_builtin_subagents_creates_agents(self, monkeypatch):
        from yuxi.services import subagent_service as service_module

        created_agents = []

        class MockRepo:
            def __init__(self, session):
                pass

            async def exists_name(self, name):
                return False

            async def create(self, **kwargs):
                created_agents.append(kwargs)
                return MagicMock()

        @asynccontextmanager
        async def mock_session_context(*args, **kwargs):
            yield MagicMock()

        class MockPgManager:
            get_async_session_context = mock_session_context

        monkeypatch.setattr(service_module, "SubAgentRepository", MockRepo)
        monkeypatch.setattr(service_module, "pg_manager", MockPgManager())

        await service_module.init_builtin_subagents()

        assert len(created_agents) == 2
        agent_names = [a["name"] for a in created_agents]
        assert "research-agent" in agent_names
        assert "critique-agent" in agent_names

    @pytest.mark.asyncio
    async def test_get_subagent_specs_returns_list(self, monkeypatch):
        from yuxi.services import subagent_service as service_module

        mock_spec = {
            "name": "test-agent",
            "description": "Test",
            "system_prompt": "You are a test",
            "tools": ["tool_a"],
        }

        class MockRepo:
            def __init__(self, session):
                pass

            async def list_all_specs(self):
                return [mock_spec]

        @asynccontextmanager
        async def mock_session_context(*args, **kwargs):
            yield MagicMock()

        class MockPgManager:
            get_async_session_context = mock_session_context

        monkeypatch.setattr(service_module, "SubAgentRepository", MockRepo)
        monkeypatch.setattr(service_module, "pg_manager", MockPgManager())

        result = await service_module.get_subagent_specs()

        assert len(result) == 1
        assert result[0]["name"] == "test-agent"

    @pytest.mark.asyncio
    async def test_get_subagent_specs_returns_defensive_copy(self, monkeypatch):
        from yuxi.services import subagent_service as service_module

        service_module._subagent_specs_cache = [
            {
                "name": "test-agent",
                "description": "Test",
                "system_prompt": "You are a test",
                "tools": ["tool_a"],
            }
        ]

        first = await service_module.get_subagent_specs()
        first[0]["tools"].append("tool_b")
        second = await service_module.get_subagent_specs()

        assert second[0]["tools"] == ["tool_a"]
        service_module.clear_specs_cache()

# =============================================================================
# Model Tests
# =============================================================================

class TestSubAgentModel:
    def test_to_dict(self):
        now = utc_now_naive()
        agent = SubAgent(
            name="test-agent",
            description="Test agent",
            system_prompt="You are a test",
            tools=["tool_a", "tool_b"],
            model="gpt-4",
            is_builtin=False,
            created_by="admin",
            updated_by="admin",
            created_at=now,
            updated_at=now,
        )

        result = agent.to_dict()

        assert result["name"] == "test-agent"
        assert result["description"] == "Test agent"
        assert result["system_prompt"] == "You are a test"
        assert result["tools"] == ["tool_a", "tool_b"]
        assert result["model"] == "gpt-4"
        assert result["is_builtin"] is False
        assert result["created_by"] == "admin"

    def test_to_subagent_spec(self):
        agent = SubAgent(
            name="test-agent",
            description="Test agent",
            system_prompt="You are a test",
            tools=["tool_a"],
            model="gpt-4",
            is_builtin=False,
            created_by="admin",
            updated_by="admin",
            created_at=utc_now_naive(),
            updated_at=utc_now_naive(),
        )

        spec = agent.to_subagent_spec()

        assert spec["name"] == "test-agent"
        assert spec["description"] == "Test agent"
        assert spec["system_prompt"] == "You are a test"
        assert spec["tools"] == ["tool_a"]
        assert spec["model"] == "gpt-4"

    def test_to_subagent_spec_no_model(self):
        agent = SubAgent(
            name="test-agent",
            description="Test agent",
            system_prompt="You are a test",
            tools=[],
            model=None,
            is_builtin=False,
            created_by="admin",
            updated_by="admin",
            created_at=utc_now_naive(),
            updated_at=utc_now_naive(),
        )

        spec = agent.to_subagent_spec()

        assert "model" not in spec


class TestDeepAgentSubagentSelection:
    @pytest.mark.asyncio
    async def test_get_subagents_from_names_filters_and_resolves_tools(self, monkeypatch):
        from yuxi.services import subagent_service as service_module

        async def fake_get_specs(_db=None):
            return [
                {
                    "name": "research-agent",
                    "description": "r",
                    "system_prompt": "s",
                    "tools": ["tool_a"],
                },
                {
                    "name": "critique-agent",
                    "description": "c",
                    "system_prompt": "s",
                    "tools": [],
                },
            ]

        mock_tool = MagicMock()
        mock_tool.name = "tool_a"

        monkeypatch.setattr(service_module, "get_subagent_specs", fake_get_specs)
        monkeypatch.setattr("yuxi.agents.toolkits.get_all_tool_instances", lambda: [mock_tool])

        resolved_specs = await service_module.get_subagents_from_names(["research-agent", "missing-agent"])

        assert [item["name"] for item in resolved_specs] == ["research-agent"]
        assert resolved_specs[0]["tools"] == [mock_tool]