tests/test_admin_management_api.py

from __future__ import annotations

from pathlib import Path

from fastapi.testclient import TestClient

from zai2api.account_pool import AccountPool
from zai2api.config import Settings
from zai2api.db import Database
from zai2api.server import create_app
from zai2api.zai_client import SessionState, UpstreamResult


class FakeManagedClient:
    def __init__(self, *, identity: str):
        self.identity = identity

    async def ensure_session(self, force_refresh: bool = False) -> SessionState:
        return SessionState(
            token=f"session-{self.identity}",
            user_id=f"user-{self.identity}",
            name=self.identity,
            email=f"{self.identity}@example.com",
            role="user",
        )

    async def verify_completion_version(self) -> int:
        return 2

    async def collect_prompt(
        self,
        *,
        prompt: str,
        model: str,
        enable_thinking: bool,
        auto_web_search: bool,
    ) -> UpstreamResult:
        return UpstreamResult(
            answer_text=f"ok:{prompt}",
            reasoning_text="reasoning",
            usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
            finish_reason="stop",
        )

    async def stream_prompt(self, **_: object):
        if False:
            yield None

    async def aclose(self) -> None:
        return None


def make_settings(tmp_path: Path) -> Settings:
    return Settings(
        host="127.0.0.1",
        port=8000,
        log_level="info",
        zai_base_url="https://chat.z.ai",
        zai_jwt=None,
        zai_session_token=None,
        default_model="glm-5",
        request_timeout=120.0,
        database_path=str(tmp_path / "state.db"),
        panel_password_env=None,
        api_password_env=None,
        admin_cookie_name="zai2api_admin_session",
        admin_session_ttl_hours=24,
        admin_cookie_secure=False,
        account_poll_interval_seconds=0,
    )


def build_client(tmp_path: Path) -> TestClient:
    settings = make_settings(tmp_path)
    db = Database(settings.database_path)
    pool = AccountPool(
        settings,
        db,
        client_factory=lambda jwt, session_token: FakeManagedClient(
            identity=(jwt or session_token or "fallback").replace("jwt-", "")
        ),
    )
    app = create_app(settings, account_pool=pool)
    return TestClient(app)


def login(client: TestClient) -> None:
    response = client.post("/api/admin/login", json={"password": "123456"})
    assert response.status_code == 200


def test_admin_can_add_list_and_toggle_accounts(tmp_path: Path) -> None:
    with build_client(tmp_path) as client:
        login(client)

        added = client.post("/api/admin/accounts", json={"jwt": "jwt-alpha"})
        assert added.status_code == 200
        account = added.json()["account"]
        assert account["email"] == "alpha@example.com"
        assert account["masked_jwt"] != "jwt-alpha"
        assert account["masked_jwt"] is not None

        listed = client.get("/api/admin/accounts")
        assert listed.status_code == 200
        assert len(listed.json()["accounts"]) == 1

        disabled = client.post(f"/api/admin/accounts/{account['id']}/disable")
        assert disabled.status_code == 200
        assert disabled.json()["account"]["enabled"] is False

        checked = client.post(f"/api/admin/accounts/{account['id']}/check")
        assert checked.status_code == 200
        assert checked.json()["account"]["status"] == "active"


def test_admin_can_update_security_settings_and_read_logs(tmp_path: Path) -> None:
    with build_client(tmp_path) as client:
        login(client)

        updated = client.post(
            "/api/admin/settings/security",
            json={
                "panel_password": "new-panel",
                "api_password": "api-secret",
                "log_retention_days": 14,
            },
        )
        assert updated.status_code == 200
        payload = updated.json()
        assert payload["panel_password"]["source"] == "database"
        assert payload["api_password"]["enabled"] is True
        assert payload["log_retention"]["days"] == 14
        assert payload["log_retention"]["source"] == "database"

        client.post("/api/admin/logout")
        relogin = client.post("/api/admin/login", json={"password": "new-panel"})
        assert relogin.status_code == 200

        logs = client.get("/api/admin/logs?limit=20")
        assert logs.status_code == 200
        messages = [item["message"] for item in logs.json()["logs"]]
        assert "Updated security settings" in messages


def test_openai_requests_are_written_to_audit_logs(tmp_path: Path) -> None:
    with build_client(tmp_path) as client:
        login(client)
        added = client.post("/api/admin/accounts", json={"jwt": "jwt-alpha"})
        assert added.status_code == 200

        models = client.get("/v1/models")
        assert models.status_code == 200

        completion = client.post(
            "/v1/chat/completions",
            json={"model": "glm-5", "messages": [{"role": "user", "content": "hello"}]},
        )
        assert completion.status_code == 200

        response_api = client.post(
            "/v1/responses",
            json={"model": "glm-5", "input": "hello"},
        )
        assert response_api.status_code == 200

        logs = client.get("/api/admin/logs?limit=50")
        assert logs.status_code == 200
        messages = [item["message"] for item in logs.json()["logs"]]
        assert "Listed available models" in messages
        assert "Completed chat completion request" in messages
        assert "Completed responses request" in messages