import importlib import os import sys from pathlib import Path from fastapi.testclient import TestClient os.environ["JOB_BACKEND"] = "memory" BACKEND_DIR = Path(__file__).resolve().parents[1] 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 prompt # type: ignore main = importlib.import_module("main") HEADERS = {"X-API-Key": main.API_KEY} def setup_function(): job_system.reset_job_manager() main._handlers_registered = False def _payload(): return { "prefix": "Before", "suffix": "After", "languageId": "markdown", "instruction": "expand", "pro_thinking": "medium", "privacy_mode": True, } def test_pro_queue_full_returns_429(monkeypatch): async def fake_queue_job(*args, **kwargs): raise job_system.QueueFullError("pro_completion queue is full") monkeypatch.setattr(main, "_queue_job", fake_queue_job) with TestClient(main.app) as client: response = client.post("/v1/pro/completions", headers=HEADERS, json=_payload()) assert response.status_code == 429 def test_pro_status_missing_returns_404(): with TestClient(main.app) as client: response = client.get("/v1/pro/completions/status/missing", headers=HEADERS) assert response.status_code == 404 def test_pro_prompt_uses_pro_specific_instruction(): system_prompt, user_prompt = prompt.build_pro_completion_prompts( prefix="欢迎使用 LLM-IN-TEXT\n\n即时可用的 LLM 系统", suffix="", language_id="markdown", instruction="", pro_thinking_level="high", ) combined = f"{system_prompt}\n{user_prompt}".lower() assert "[pro] model for llm-in-text" in combined assert "pro_mode: true" in combined assert "pro_thinking_level: high" in combined def test_pro_stream_returns_standard_events(monkeypatch): async def fake_stream_events(*args, **kwargs): yield "thinking", "" yield "content", "深度" yield "content", "回答" monkeypatch.setattr(job_handlers, "stream_ollama_events", fake_stream_events) with TestClient(main.app) as client: with client.stream("POST", "/v1/pro/completions", headers=HEADERS, json=_payload()) as resp: assert resp.status_code == 200 body = "".join(resp.iter_text()) assert "event: queued" in body assert "event: started" in body assert "event: progress" in body assert "event: result" in body assert "event: done" in body assert "深度" in body assert "回答" in body