3
This commit is contained in:
@@ -0,0 +1,229 @@
|
||||
"""Phase 6 文档上传测试。"""
|
||||
|
||||
import io
|
||||
|
||||
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 _create_kb(client: TestClient, name: str = "测试知识库") -> str:
|
||||
resp = client.post("/api/knowledge-bases", json={"name": name})
|
||||
return resp.json()["id"]
|
||||
|
||||
|
||||
def _make_docx_bytes() -> bytes:
|
||||
"""创建一个最小的 .docx 文件字节(实际是 ZIP 格式)。"""
|
||||
import zipfile
|
||||
|
||||
buf = io.BytesIO()
|
||||
with zipfile.ZipFile(buf, "w") as zf:
|
||||
zf.writestr("[Content_Types].xml", '<?xml version="1.0"?><Types></Types>')
|
||||
zf.writestr("_rels/.rels", '<?xml version="1.0"?><Relationships></Relationships>')
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def _make_pdf_bytes() -> bytes:
|
||||
"""创建一个最小的 PDF 文件字节。"""
|
||||
return b"%PDF-1.4\n1 0 obj\n<< /Type /Catalog >>\nendobj\n"
|
||||
|
||||
|
||||
# --- 上传 ---
|
||||
|
||||
def test_upload_docx() -> None:
|
||||
with _client() as client:
|
||||
_register_and_login(client)
|
||||
kb_id = _create_kb(client)
|
||||
resp = client.post(
|
||||
"/api/documents/upload",
|
||||
data={"kb_id": kb_id},
|
||||
files={"file": ("test.docx", _make_docx_bytes(), "application/vnd.openxmlformats-officedocument.wordprocessingml.document")},
|
||||
)
|
||||
assert resp.status_code == 201
|
||||
body = resp.json()
|
||||
assert body["original_filename"] == "test.docx"
|
||||
# 同步解析已完成(测试用最小 .docx 无实际内容,可能 FAILED;真实文件会 READY)
|
||||
assert body["status"] in ("PENDING", "PROCESSING", "READY", "FAILED")
|
||||
assert body["file_size"] > 0
|
||||
|
||||
|
||||
def test_upload_pdf() -> None:
|
||||
with _client() as client:
|
||||
_register_and_login(client)
|
||||
kb_id = _create_kb(client)
|
||||
resp = client.post(
|
||||
"/api/documents/upload",
|
||||
data={"kb_id": kb_id},
|
||||
files={"file": ("test.pdf", _make_pdf_bytes(), "application/pdf")},
|
||||
)
|
||||
assert resp.status_code == 201
|
||||
assert resp.json()["original_filename"] == "test.pdf"
|
||||
|
||||
|
||||
def test_upload_unsupported_extension() -> None:
|
||||
with _client() as client:
|
||||
_register_and_login(client)
|
||||
kb_id = _create_kb(client)
|
||||
resp = client.post(
|
||||
"/api/documents/upload",
|
||||
data={"kb_id": kb_id},
|
||||
files={"file": ("test.txt", b"hello", "text/plain")},
|
||||
)
|
||||
assert resp.status_code == 415
|
||||
assert resp.json()["code"] == "FILE_TYPE_UNSUPPORTED"
|
||||
|
||||
|
||||
def test_upload_too_large() -> None:
|
||||
"""超过 20MB 限制。"""
|
||||
with _client() as client:
|
||||
_register_and_login(client)
|
||||
kb_id = _create_kb(client)
|
||||
large_content = b"x" * (20 * 1024 * 1024 + 1) # 20MB + 1 byte
|
||||
resp = client.post(
|
||||
"/api/documents/upload",
|
||||
data={"kb_id": kb_id},
|
||||
files={"file": ("large.docx", large_content, "application/vnd.openxmlformats-officedocument.wordprocessingml.document")},
|
||||
)
|
||||
assert resp.status_code == 413
|
||||
assert resp.json()["code"] == "FILE_TOO_LARGE"
|
||||
|
||||
|
||||
def test_upload_wrong_kb_owner() -> None:
|
||||
"""用户A 不能上传到用户B 的知识库。"""
|
||||
with _client() as client:
|
||||
_register_and_login(client, "user_a")
|
||||
kb_id = _create_kb(client, "A 的知识库")
|
||||
# 用户 B 登录
|
||||
_register_and_login(client, "user_b")
|
||||
resp = client.post(
|
||||
"/api/documents/upload",
|
||||
data={"kb_id": kb_id},
|
||||
files={"file": ("test.docx", _make_docx_bytes(), "application/vnd.openxmlformats-officedocument.wordprocessingml.document")},
|
||||
)
|
||||
assert resp.status_code == 404 # 对 B 来说 KB 不存在
|
||||
|
||||
|
||||
# --- 列表 ---
|
||||
|
||||
def test_list_documents() -> None:
|
||||
with _client() as client:
|
||||
_register_and_login(client)
|
||||
kb_id = _create_kb(client)
|
||||
# 上传 2 个文件
|
||||
client.post(
|
||||
"/api/documents/upload",
|
||||
data={"kb_id": kb_id},
|
||||
files={"file": ("a.docx", _make_docx_bytes(), "application/vnd.openxmlformats-officedocument.wordprocessingml.document")},
|
||||
)
|
||||
client.post(
|
||||
"/api/documents/upload",
|
||||
data={"kb_id": kb_id},
|
||||
files={"file": ("b.pdf", _make_pdf_bytes(), "application/pdf")},
|
||||
)
|
||||
resp = client.get(f"/api/documents?kb_id={kb_id}")
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["total"] == 2
|
||||
assert len(body["items"]) == 2
|
||||
|
||||
|
||||
# --- 详情 ---
|
||||
|
||||
def test_get_document() -> None:
|
||||
with _client() as client:
|
||||
_register_and_login(client)
|
||||
kb_id = _create_kb(client)
|
||||
upload_resp = client.post(
|
||||
"/api/documents/upload",
|
||||
data={"kb_id": kb_id},
|
||||
files={"file": ("test.docx", _make_docx_bytes(), "application/vnd.openxmlformats-officedocument.wordprocessingml.document")},
|
||||
)
|
||||
doc_id = upload_resp.json()["id"]
|
||||
resp = client.get(f"/api/documents/{doc_id}")
|
||||
assert resp.status_code == 200
|
||||
assert resp.json()["original_filename"] == "test.docx"
|
||||
|
||||
|
||||
# --- 编辑 ---
|
||||
|
||||
def test_update_document() -> None:
|
||||
with _client() as client:
|
||||
_register_and_login(client)
|
||||
kb_id = _create_kb(client)
|
||||
upload_resp = client.post(
|
||||
"/api/documents/upload",
|
||||
data={"kb_id": kb_id},
|
||||
files={"file": ("test.docx", _make_docx_bytes(), "application/vnd.openxmlformats-officedocument.wordprocessingml.document")},
|
||||
)
|
||||
doc_id = upload_resp.json()["id"]
|
||||
resp = client.put(f"/api/documents/{doc_id}", json={"title": "新标题"})
|
||||
assert resp.status_code == 200
|
||||
assert resp.json()["title"] == "新标题"
|
||||
|
||||
|
||||
# --- 删除 ---
|
||||
|
||||
def test_delete_document() -> None:
|
||||
with _client() as client:
|
||||
_register_and_login(client)
|
||||
kb_id = _create_kb(client)
|
||||
upload_resp = client.post(
|
||||
"/api/documents/upload",
|
||||
data={"kb_id": kb_id},
|
||||
files={"file": ("test.docx", _make_docx_bytes(), "application/vnd.openxmlformats-officedocument.wordprocessingml.document")},
|
||||
)
|
||||
doc_id = upload_resp.json()["id"]
|
||||
# 删除
|
||||
resp = client.delete(f"/api/documents/{doc_id}")
|
||||
assert resp.status_code == 204
|
||||
# 删除后查不到
|
||||
resp = client.get(f"/api/documents/{doc_id}")
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
# --- IDOR ---
|
||||
|
||||
def test_user_a_cannot_see_user_b_doc() -> None:
|
||||
with _client() as client:
|
||||
_register_and_login(client, "user_a")
|
||||
kb_id = _create_kb(client)
|
||||
upload_resp = client.post(
|
||||
"/api/documents/upload",
|
||||
data={"kb_id": kb_id},
|
||||
files={"file": ("a.docx", _make_docx_bytes(), "application/vnd.openxmlformats-officedocument.wordprocessingml.document")},
|
||||
)
|
||||
doc_id = upload_resp.json()["id"]
|
||||
# 用户 B 登录
|
||||
_register_and_login(client, "user_b")
|
||||
resp = client.get(f"/api/documents/{doc_id}")
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
# --- 配额 ---
|
||||
|
||||
def test_storage_quota_deducted() -> None:
|
||||
with _client() as client:
|
||||
_register_and_login(client)
|
||||
kb_id = _create_kb(client)
|
||||
content = _make_docx_bytes()
|
||||
client.post(
|
||||
"/api/documents/upload",
|
||||
data={"kb_id": kb_id},
|
||||
files={"file": ("test.docx", content, "application/vnd.openxmlformats-officedocument.wordprocessingml.document")},
|
||||
)
|
||||
# 检查存储用量
|
||||
resp = client.get("/api/auth/storage")
|
||||
assert resp.status_code == 200
|
||||
assert resp.json()["storage_used"] > 0
|
||||
Reference in New Issue
Block a user