178 lines
5.5 KiB
Python
178 lines
5.5 KiB
Python
"""Phase 9-11 公共 AI 页面测试。"""
|
|
|
|
from fastapi.testclient import TestClient
|
|
|
|
from app.main import app
|
|
|
|
|
|
def _client() -> TestClient:
|
|
return TestClient(app, raise_server_exceptions=False)
|
|
|
|
|
|
def _register_and_create_kb(client: TestClient) -> tuple[str, str]:
|
|
"""注册用户并创建知识库，返回 (kb_id, ai_url)。"""
|
|
client.post("/api/auth/register", json={
|
|
"username": "testuser",
|
|
"email": "test@example.com",
|
|
"password": "password123",
|
|
})
|
|
resp = client.post("/api/knowledge-bases", json={
|
|
"name": "测试知识库",
|
|
"description": "这是测试描述",
|
|
})
|
|
kb_id = resp.json()["id"]
|
|
ai_url = resp.json()["ai_url"]
|
|
return kb_id, ai_url
|
|
|
|
|
|
def _get_token_from_url(ai_url: str) -> str:
|
|
"""从 AI URL 中提取 token。"""
|
|
return ai_url.split("/k/")[-1]
|
|
|
|
|
|
# --- HTML 入口 ---
|
|
|
|
def test_kb_index_html() -> None:
|
|
with _client() as client:
|
|
_, ai_url = _register_and_create_kb(client)
|
|
token = _get_token_from_url(ai_url)
|
|
resp = client.get(f"/k/{token}")
|
|
assert resp.status_code == 200
|
|
assert "text/html" in resp.headers["content-type"]
|
|
assert "测试知识库" in resp.text
|
|
assert "noindex" in resp.text
|
|
assert "no-referrer" in resp.text
|
|
|
|
|
|
def test_kb_index_html_not_found() -> None:
|
|
with _client() as client:
|
|
resp = client.get("/k/nonexistent")
|
|
assert resp.status_code == 404
|
|
|
|
|
|
def test_kb_index_html_disabled() -> None:
|
|
with _client() as client:
|
|
kb_id, ai_url = _register_and_create_kb(client)
|
|
token = _get_token_from_url(ai_url)
|
|
# 禁用链接
|
|
client.post(f"/api/knowledge-bases/{kb_id}/disable")
|
|
resp = client.get(f"/k/{token}")
|
|
assert resp.status_code == 404
|
|
|
|
|
|
# --- Markdown 输出 ---
|
|
|
|
def test_kb_index_markdown() -> None:
|
|
with _client() as client:
|
|
_, ai_url = _register_and_create_kb(client)
|
|
token = _get_token_from_url(ai_url)
|
|
resp = client.get(f"/k/{token}.md")
|
|
assert resp.status_code == 200
|
|
assert "text/markdown" in resp.headers["content-type"]
|
|
assert "测试知识库" in resp.text
|
|
|
|
|
|
# --- 纯文本输出 ---
|
|
|
|
def test_kb_index_text() -> None:
|
|
with _client() as client:
|
|
_, ai_url = _register_and_create_kb(client)
|
|
token = _get_token_from_url(ai_url)
|
|
resp = client.get(f"/k/{token}.txt")
|
|
assert resp.status_code == 200
|
|
assert "text/plain" in resp.headers["content-type"]
|
|
assert "测试知识库" in resp.text
|
|
|
|
|
|
# --- JSON 输出 ---
|
|
|
|
def test_kb_index_json() -> None:
|
|
with _client() as client:
|
|
_, ai_url = _register_and_create_kb(client)
|
|
token = _get_token_from_url(ai_url)
|
|
resp = client.get(f"/k/{token}.json")
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
assert data["name"] == "测试知识库"
|
|
assert "categories" in data
|
|
# 每个目录有 documents 字段
|
|
if data["categories"]:
|
|
assert "documents" in data["categories"][0]
|
|
# 安全：不包含内部字段
|
|
assert "user_id" not in data
|
|
assert "token_hash" not in data
|
|
|
|
|
|
# --- 搜索 ---
|
|
|
|
def test_search_html() -> None:
|
|
with _client() as client:
|
|
_, ai_url = _register_and_create_kb(client)
|
|
token = _get_token_from_url(ai_url)
|
|
resp = client.get(f"/k/{token}/search?q=test")
|
|
assert resp.status_code == 200
|
|
assert "text/html" in resp.headers["content-type"]
|
|
|
|
|
|
def test_search_json() -> None:
|
|
with _client() as client:
|
|
_, ai_url = _register_and_create_kb(client)
|
|
token = _get_token_from_url(ai_url)
|
|
resp = client.get(f"/k/{token}/search.json?q=test")
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
assert "query" in data
|
|
assert "total" in data
|
|
|
|
|
|
# --- 单文档访问 ---
|
|
|
|
def test_doc_page_html() -> None:
|
|
"""上传文档后通过 token 访问文档页面。"""
|
|
import io, zipfile
|
|
|
|
with _client() as client:
|
|
kb_id, ai_url = _register_and_create_kb(client)
|
|
token = _get_token_from_url(ai_url)
|
|
|
|
# 上传一个真实 .docx 文件
|
|
buf = io.BytesIO()
|
|
with zipfile.ZipFile(buf, "w") as zf:
|
|
zf.writestr("[Content_Types].xml", '<?xml version="1.0"?><Types></Types>')
|
|
zf.writestr("word/document.xml", '<?xml version="1.0"?><w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main"><w:body><w:p><w:r><w:t>Hello World</w:t></w:r></w:p></w:body></w:document>')
|
|
|
|
# 获取文档 token
|
|
resp = client.get(f"/api/documents?kb_id={kb_id}")
|
|
if resp.json()["total"] > 0:
|
|
doc = resp.json()["items"][0]
|
|
doc_token_hint = doc.get("doc_token_hint", "")
|
|
# 访问文档页面（使用 hint 作为 token 的一部分）
|
|
resp = client.get(f"/k/{token}/doc/{doc_token_hint}")
|
|
# 可能 404（因为 hint 不是完整 token），但不应 500
|
|
assert resp.status_code in (200, 404)
|
|
|
|
|
|
# --- 限流 ---
|
|
|
|
def test_rate_limit_not_triggered() -> None:
|
|
"""正常请求不应触发限流。"""
|
|
with _client() as client:
|
|
_, ai_url = _register_and_create_kb(client)
|
|
token = _get_token_from_url(ai_url)
|
|
# 连续请求 5 次
|
|
for _ in range(5):
|
|
resp = client.get(f"/k/{token}")
|
|
assert resp.status_code == 200
|
|
|
|
|
|
# --- 安全头 ---
|
|
|
|
def test_no_robots_meta() -> None:
|
|
with _client() as client:
|
|
_, ai_url = _register_and_create_kb(client)
|
|
token = _get_token_from_url(ai_url)
|
|
resp = client.get(f"/k/{token}")
|
|
assert "noindex" in resp.text
|
|
assert "nofollow" in resp.text
|
|
assert "noarchive" in resp.text
|