amb_rag/backend/tests/test_knowledge_bases.py

"""Phase 4 知识库 CRUD + Token 管理测试。"""

from fastapi.testclient import TestClient

from app.main import app


def _client() -> TestClient:
    return TestClient(app, raise_server_exceptions=False)


def _register_and_login(client: TestClient, username: str = "testuser") -> None:
    client.post("/api/auth/register", json={
        "username": username,
        "email": f"{username}@example.com",
        "password": "password123",
    })


# --- 创建 ---

def test_create_knowledge_base() -> None:
    with _client() as client:
        _register_and_login(client)
        resp = client.post("/api/knowledge-bases", json={
            "name": "公司知识库",
            "description": "公司相关文档",
        })
    assert resp.status_code == 201
    body = resp.json()
    assert body["name"] == "公司知识库"
    assert body["description"] == "公司相关文档"
    assert body["enabled"] is True
    assert body["token_hint"] is not None
    assert body["ai_url"].startswith("/k/")


def test_create_kb_name_required() -> None:
    with _client() as client:
        _register_and_login(client)
        resp = client.post("/api/knowledge-bases", json={"name": ""})
    assert resp.status_code == 422


# --- 列表 ---

def test_list_knowledge_bases() -> None:
    with _client() as client:
        _register_and_login(client)
        client.post("/api/knowledge-bases", json={"name": "KB1"})
        client.post("/api/knowledge-bases", json={"name": "KB2"})
        resp = client.get("/api/knowledge-bases")
    assert resp.status_code == 200
    body = resp.json()
    assert body["total"] == 2
    assert len(body["items"]) == 2


def test_list_kb_pagination() -> None:
    with _client() as client:
        _register_and_login(client)
        for i in range(5):
            client.post("/api/knowledge-bases", json={"name": f"KB{i}"})
        resp = client.get("/api/knowledge-bases?page=1&page_size=2")
    assert resp.status_code == 200
    body = resp.json()
    assert body["total"] == 5
    assert len(body["items"]) == 2


# --- 详情 ---

def test_get_knowledge_base() -> None:
    with _client() as client:
        _register_and_login(client)
        create_resp = client.post("/api/knowledge-bases", json={"name": "详情测试"})
        kb_id = create_resp.json()["id"]
        resp = client.get(f"/api/knowledge-bases/{kb_id}")
    assert resp.status_code == 200
    assert resp.json()["name"] == "详情测试"


def test_get_kb_not_found() -> None:
    with _client() as client:
        _register_and_login(client)
        resp = client.get("/api/knowledge-bases/nonexistent")
    assert resp.status_code == 404


# --- 编辑 ---

def test_update_knowledge_base() -> None:
    with _client() as client:
        _register_and_login(client)
        create_resp = client.post("/api/knowledge-bases", json={"name": "旧名称"})
        kb_id = create_resp.json()["id"]
        resp = client.put(f"/api/knowledge-bases/{kb_id}", json={"name": "新名称"})
    assert resp.status_code == 200
    assert resp.json()["name"] == "新名称"


# --- 删除 ---

def test_delete_knowledge_base() -> None:
    with _client() as client:
        _register_and_login(client)
        create_resp = client.post("/api/knowledge-bases", json={"name": "待删除"})
        kb_id = create_resp.json()["id"]
        # 删除
        resp = client.delete(f"/api/knowledge-bases/{kb_id}")
        assert resp.status_code == 204
        # 删除后对自己也查不到（status=DELETED）
        resp = client.get(f"/api/knowledge-bases/{kb_id}")
    assert resp.status_code == 404


# --- Token 管理 ---

def test_regenerate_token() -> None:
    with _client() as client:
        _register_and_login(client)
        create_resp = client.post("/api/knowledge-bases", json={"name": "Token 测试"})
        kb_id = create_resp.json()["id"]
        old_hint = create_resp.json()["token_hint"]
        resp = client.post(f"/api/knowledge-bases/{kb_id}/regenerate-token")
    assert resp.status_code == 200
    body = resp.json()
    assert body["ai_url"].startswith("/k/")
    assert body["token"] != ""
    # hint 可能相同（概率极低）但 token 应该不同


def test_enable_disable() -> None:
    with _client() as client:
        _register_and_login(client)
        create_resp = client.post("/api/knowledge-bases", json={"name": "开关测试"})
        kb_id = create_resp.json()["id"]
        # 禁用
        resp = client.post(f"/api/knowledge-bases/{kb_id}/disable")
        assert resp.status_code == 200
        assert resp.json()["enabled"] is False
        # 启用
        resp = client.post(f"/api/knowledge-bases/{kb_id}/enable")
        assert resp.status_code == 200
        assert resp.json()["enabled"] is True


def test_get_link() -> None:
    with _client() as client:
        _register_and_login(client)
        create_resp = client.post("/api/knowledge-bases", json={"name": "链接测试"})
        kb_id = create_resp.json()["id"]
        resp = client.get(f"/api/knowledge-bases/{kb_id}/link")
    assert resp.status_code == 200
    body = resp.json()
    assert body["ai_url"].startswith("/k/")
    assert len(body["token"]) > 0


# --- IDOR 防护 ---

def test_user_a_cannot_access_user_b_kb() -> None:
    with _client() as client:
        # 用户 A 创建知识库
        _register_and_login(client, "user_a")
        create_resp = client.post("/api/knowledge-bases", json={"name": "A 的知识库"})
        kb_id = create_resp.json()["id"]
        # 用户 B 登录
        _register_and_login(client, "user_b")
        resp = client.get(f"/api/knowledge-bases/{kb_id}")
    assert resp.status_code == 404  # 对 B 来说"不存在"