"""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", '') zf.writestr("_rels/.rels", '') 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