KiloStar/tests/unit/test_api_agent_template.py

"""``api/agent.py`` persona template 路由：CRUD 鉴权与 node_affinity 校验。"""

from __future__ import annotations

import types
from unittest.mock import AsyncMock

import pytest
from fastapi import FastAPI
from httpx import AsyncClient, ASGITransport

from kilostar.api.agent import agent_router, _VALID_AFFINITIES
from kilostar.utils.access import Accessor, TokenData
from kilostar.core.postgres_database.model import UserAuthority


def _fake_user(user_id: str = "alice"):
    return TokenData(user_id=user_id, username=user_id)


def _tpl(owner: str = "alice", is_builtin: bool = False):
    return types.SimpleNamespace(
        template_id="tpl1", name="MyBot", agent_type="ordinary",
        description="d", system_prompt="s", provider_title="openai",
        model_id="gpt-4", tools=[], tags=[], is_builtin=is_builtin, owner_id=owner,
    )


@pytest.fixture
def app(monkeypatch):
    import kilostar.utils.check_user.role_check as rc
    monkeypatch.setattr(rc, "get_authority", AsyncMock(return_value=UserAuthority.USER))
    _app = FastAPI()
    _app.include_router(agent_router)
    _app.dependency_overrides[Accessor.get_current_user] = lambda: _fake_user()
    return _app


def _register_pg(fake_actors, **kwargs):
    pg = types.SimpleNamespace(**kwargs)
    fake_actors.register("postgres_database", pg)
    return pg


# ── node_affinity 校验（纯 pydantic） ─────────────────────────────────────


def test_valid_affinities_accepted():
    from kilostar.api.agent import WorkerIndividualCreate
    for aff in _VALID_AFFINITIES:
        m = WorkerIndividualCreate(
            agent_name="x", agent_type="ordinary", description="d",
            provider_title="p", model_id="m", system_prompt="s",
            output_template={}, bound_skill={}, workspace=[], node_affinity=aff,
        )
        assert m.node_affinity == aff


def test_invalid_affinity_raises():
    from pydantic import ValidationError
    from kilostar.api.agent import WorkerIndividualCreate
    with pytest.raises(ValidationError):
        WorkerIndividualCreate(
            agent_name="x", agent_type="ordinary", description="d",
            provider_title="p", model_id="m", system_prompt="s",
            output_template={}, bound_skill={}, workspace=[], node_affinity="bad",
        )


def test_update_invalid_affinity_raises():
    from pydantic import ValidationError
    from kilostar.api.agent import WorkerIndividualUpdate
    with pytest.raises(ValidationError):
        WorkerIndividualUpdate(node_affinity="bad")


def test_update_none_affinity_ok():
    from kilostar.api.agent import WorkerIndividualUpdate
    assert WorkerIndividualUpdate(node_affinity=None).node_affinity is None


# ── template API 路由 ──────────────────────────────────────────────────────


@pytest.mark.asyncio
async def test_list_templates(app, fake_actors):
    pg = types.SimpleNamespace(
        list_templates=types.SimpleNamespace(remote=AsyncMock(return_value=[]))
    )
    fake_actors.register("postgres_database", pg)
    async with AsyncClient(transport=ASGITransport(app=app), base_url="http://t") as c:
        r = await c.get("/api/v1/agent/template")
    assert r.status_code == 200
    assert r.json()["templates"] == []


@pytest.mark.asyncio
async def test_create_template(app, fake_actors):
    tpl = _tpl()
    pg = types.SimpleNamespace(
        add_template=types.SimpleNamespace(remote=AsyncMock(return_value=tpl))
    )
    fake_actors.register("postgres_database", pg)
    async with AsyncClient(transport=ASGITransport(app=app), base_url="http://t") as c:
        r = await c.post("/api/v1/agent/template", json={"name": "test"})
    assert r.status_code == 200
    assert r.json()["template_id"] == "tpl1"


@pytest.mark.asyncio
async def test_delete_template_not_found(app, fake_actors):
    pg = types.SimpleNamespace(
        get_template=types.SimpleNamespace(remote=AsyncMock(return_value=None))
    )
    fake_actors.register("postgres_database", pg)
    async with AsyncClient(transport=ASGITransport(app=app), base_url="http://t") as c:
        r = await c.delete("/api/v1/agent/template/missing")
    assert r.status_code == 404


@pytest.mark.asyncio
async def test_delete_builtin_template_forbidden(app, fake_actors):
    pg = types.SimpleNamespace(
        get_template=types.SimpleNamespace(remote=AsyncMock(return_value=_tpl(is_builtin=True)))
    )
    fake_actors.register("postgres_database", pg)
    async with AsyncClient(transport=ASGITransport(app=app), base_url="http://t") as c:
        r = await c.delete("/api/v1/agent/template/tpl1")
    assert r.status_code == 403


@pytest.mark.asyncio
async def test_delete_other_users_template_forbidden(app, fake_actors):
    pg = types.SimpleNamespace(
        get_template=types.SimpleNamespace(remote=AsyncMock(return_value=_tpl(owner="bob")))
    )
    fake_actors.register("postgres_database", pg)
    async with AsyncClient(transport=ASGITransport(app=app), base_url="http://t") as c:
        r = await c.delete("/api/v1/agent/template/tpl1")
    assert r.status_code == 403


@pytest.mark.asyncio
async def test_create_worker_from_template(app, fake_actors):
    worker = types.SimpleNamespace(agent_id="w1")
    pg = types.SimpleNamespace(
        get_template=types.SimpleNamespace(remote=AsyncMock(return_value=_tpl())),
        add_worker_individual=types.SimpleNamespace(remote=AsyncMock(return_value=worker)),
    )
    fake_actors.register("postgres_database", pg)
    async with AsyncClient(transport=ASGITransport(app=app), base_url="http://t") as c:
        r = await c.post("/api/v1/agent/worker/from-template/tpl1")
    assert r.status_code == 200
    assert r.json()["agent_id"] == "w1"


@pytest.mark.asyncio
async def test_create_worker_from_missing_template(app, fake_actors):
    pg = types.SimpleNamespace(
        get_template=types.SimpleNamespace(remote=AsyncMock(return_value=None))
    )
    fake_actors.register("postgres_database", pg)
    async with AsyncClient(transport=ASGITransport(app=app), base_url="http://t") as c:
        r = await c.post("/api/v1/agent/worker/from-template/nope")
    assert r.status_code == 404