import base64 import asyncio import importlib import os import sys from pathlib import Path from types import SimpleNamespace from fastapi.testclient import TestClient os.environ["JOB_BACKEND"] = "memory" os.environ["DOCS_BACKEND"] = "memory" CURRENT_DIR = Path(__file__).resolve().parent BACKEND_DIR = CURRENT_DIR.parent if str(BACKEND_DIR) not in sys.path: sys.path.insert(0, str(BACKEND_DIR)) import job_handlers # type: ignore import job_system # type: ignore import docs_store # type: ignore import risk_control # type: ignore import session_store # type: ignore import audit_store # type: ignore main = importlib.import_module("main") HEADERS = {"X-API-Key": main.API_KEY} def setup_function(): job_system.reset_job_manager() docs_store.reset_document_store() risk_control.reset_risk_controller() session_store.reset_session_store() audit_store.reset_audit_store() main._handlers_registered = False class DummyRequest: def __init__(self, host=None, headers=None): class Client: pass self.client = Client() if host is not None else None if self.client is not None: self.client.host = host self.headers = headers or {} def test_preview_short_text(): assert main._preview("Hello") == "Hello" def test_preview_long_text_truncated(): long_text = "a" * 100 assert main._preview(long_text) == long_text[:80] + "..." def test_sanitize_markdown_strips_image_markdown(): assert "![alt](image.png)" not in main._sanitize_converted_markdown("text ![alt](image.png)") def test_sanitize_inline_completion_strips_prefill(): assert main.sanitize_inline_completion_content("系统非常适合写作", prefill="系统") == "非常适合写作" def test_get_client_ip_header_overrides_host(): req = DummyRequest(host="1.2.3.4", headers={"X-Client-IP": "5.6.7.8"}) assert main.get_client_ip(req) == "5.6.7.8" def test_post_completions_without_api_key_uses_anonymous_session(monkeypatch): async def fake_call(*args, **kwargs): return {"content": "系统done", "think": ""} monkeypatch.setattr(job_handlers, "call_ollama", fake_call) with TestClient(main.app) as client: with client.stream("POST", "/v1/completions", json={ "prefix": "hello", "suffix": "", "languageId": "markdown", "model_thinking": "low", "privacy_mode": True, }) as resp: assert resp.status_code == 200 assert main.config.session_cookie_name in resp.cookies def test_post_completions_invalid_api_key_returns_403(): with TestClient(main.app) as client: resp = client.post( "/v1/completions", headers={"X-API-Key": "invalid-key"}, json={ "prefix": "hello", "suffix": "", "languageId": "markdown", "model_thinking": "low", "privacy_mode": True, }, ) assert resp.status_code == 403 def test_post_completions_returns_sse_done(monkeypatch): async def fake_call(*args, **kwargs): return {"content": "系统done", "think": ""} monkeypatch.setattr(job_handlers, "call_ollama", fake_call) with TestClient(main.app) as client: with client.stream("POST", "/v1/completions", headers=HEADERS, json={ "prefix": "hello", "suffix": "", "languageId": "markdown", "model_thinking": "low", "privacy_mode": True, }) as resp: assert resp.status_code == 200 body = "".join(resp.iter_text()) assert "event: queued" in body assert "event: started" in body assert "event: result" in body assert "event: done" in body def test_stream_job_emits_keepalive_during_idle(monkeypatch): async def fake_queue_job(*_args, **_kwargs): return "job-keepalive" class FakeManager: def register_handler(self, *_args, **_kwargs): return None async def stream_events(self, job_id): yield {"event": "queued", "job_id": job_id} yield {"event": "started", "job_id": job_id} await asyncio.sleep(0.03) yield {"event": "done", "job_id": job_id, "result": {"content": "ok"}} monkeypatch.setattr(main, "STREAM_HEARTBEAT_SECONDS", 0.01) monkeypatch.setattr(main, "_queue_job", fake_queue_job) monkeypatch.setattr(main, "get_job_manager", lambda: FakeManager()) with TestClient(main.app) as client: with client.stream("POST", "/v1/pro/completions", headers=HEADERS, json={ "prefix": "hello", "suffix": "", "languageId": "markdown", "instruction": "expand", "pro_thinking": "medium", "privacy_mode": True, }) as resp: assert resp.status_code == 200 body = "".join(resp.iter_text()) assert ": keepalive" in body assert "event: done" in body def test_post_ocr_mocked(monkeypatch): async def fake_ocr(*args, **kwargs): return "OCR result text" monkeypatch.setattr(job_handlers, "call_vlm_ocr", fake_ocr) img_b64 = base64.b64encode(b"pretend image data").decode() with TestClient(main.app) as client: with client.stream("POST", "/v1/ocr", headers=HEADERS, json={ "image": img_b64, "filename": "test.jpg", "language": "auto", }) as resp: assert resp.status_code == 200 body = "".join(resp.iter_text()) assert "OCR result text" in body def test_post_video_ocr_merges_ocr_and_asr(monkeypatch): async def fake_ocr(*args, **kwargs): return "画面文字" async def fake_asr(*args, **kwargs): return SimpleNamespace(text="音频转写") monkeypatch.setattr(job_handlers, "call_vlm_ocr", fake_ocr) monkeypatch.setattr(job_handlers, "generate_asr_response", fake_asr) monkeypatch.setattr(job_handlers, "extract_audio_wav_bytes", lambda _path: b"fake wav") video_b64 = base64.b64encode(b"pretend video data").decode() with TestClient(main.app) as client: with client.stream("POST", "/v1/ocr", headers=HEADERS, json={ "image": video_b64, "filename": "sample.mp4", "language": "auto", "media_type": "video", "mime_type": "video/mp4", }) as resp: assert resp.status_code == 200 body = "".join(resp.iter_text()) assert "视频画面 OCR" in body assert "视频音频 ASR" in body assert "画面文字" in body assert "音频转写" in body def test_post_convert_txt_returns_markdown(): content = base64.b64encode(b"hello world").decode() with TestClient(main.app) as client: with client.stream("POST", "/v1/convert", headers=HEADERS, json={ "file": content, "filename": "sample.txt", }) as resp: assert resp.status_code == 200 body = "".join(resp.iter_text()) assert "hello world" in body def test_post_convert_unsupported_extension_returns_500(): content = base64.b64encode(b"data").decode() with TestClient(main.app) as client: resp = client.post("/v1/convert", headers=HEADERS, json={ "file": content, "filename": "sample.xlsx", }) assert resp.status_code == 500 assert "仅支持" in resp.json()["error"] def test_docs_nodes_crud_round_trip(): with TestClient(main.app) as client: folder_resp = client.post("/v1/docs/folders", headers=HEADERS, json={ "name": "项目资料", "parentId": None, }) assert folder_resp.status_code == 200 folder = folder_resp.json()["node"] file_resp = client.post("/v1/docs/files/text", headers=HEADERS, json={ "name": "notes.md", "parentId": folder["id"], "content": "# hello", }) assert file_resp.status_code == 200 file_node = file_resp.json()["node"] assert file_node["previewText"] == "# hello" list_resp = client.get("/v1/docs/nodes", headers=HEADERS) assert list_resp.status_code == 200 nodes = list_resp.json()["nodes"] assert len(nodes) == 2 rename_resp = client.patch(f"/v1/docs/nodes/{file_node['id']}", headers=HEADERS, json={ "name": "renamed.md", }) assert rename_resp.status_code == 200 assert rename_resp.json()["node"]["name"] == "renamed.md" blob_resp = client.get(f"/v1/docs/files/{file_node['id']}/blob", headers=HEADERS) assert blob_resp.status_code == 200 assert blob_resp.content == b"# hello" delete_resp = client.delete(f"/v1/docs/nodes/{folder['id']}", headers=HEADERS) assert delete_resp.status_code == 200 final_list = client.get("/v1/docs/nodes", headers=HEADERS) assert final_list.status_code == 200 assert final_list.json()["nodes"] == [] def test_docs_file_upload_and_blob_replace(): with TestClient(main.app) as client: upload_resp = client.post( "/v1/docs/files/upload", headers=HEADERS, files={"file": ("image.png", b"png-bytes", "image/png")}, data={"parent_id": ""}, ) assert upload_resp.status_code == 200 node = upload_resp.json()["node"] assert node["storageKind"] == "blob" replace_resp = client.put( f"/v1/docs/files/{node['id']}/blob", headers=HEADERS, files={"file": ("photo.jpg", b"jpeg-bytes", "image/jpeg")}, ) assert replace_resp.status_code == 200 assert replace_resp.json()["node"]["name"] == "photo.jpg" blob_resp = client.get(f"/v1/docs/files/{node['id']}/blob", headers=HEADERS) assert blob_resp.status_code == 200 assert blob_resp.content == b"jpeg-bytes"