Migrate backend jobs to Redis Streams

backend/tests/test_pro_completions.py

@@ -1,28 +1,30 @@
+import importlib
 import os
 import sys
-import types
-import asyncio
-import threading
+from pathlib import Path
+
 from fastapi.testclient import TestClient
 
+os.environ["JOB_BACKEND"] = "memory"
 
-CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
-BACKEND_DIR = os.path.abspath(os.path.join(CURRENT_DIR, ".."))
-if BACKEND_DIR not in sys.path:
-    sys.path.insert(0, BACKEND_DIR)
+BACKEND_DIR = Path(__file__).resolve().parents[1]
+if str(BACKEND_DIR) not in sys.path:
+    sys.path.insert(0, str(BACKEND_DIR))
 
-if "tts_asr" not in sys.modules:
-    fake_tts_asr = types.ModuleType("tts_asr")
-    fake_tts_asr.register_tts_asr_routes = lambda app: None
-    sys.modules["tts_asr"] = fake_tts_asr
-
-import main  # type: ignore
-import pro_completions  # type: ignore
+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",
@@ -34,84 +36,52 @@ def _payload():
     }
 
 
-def setup_function():
-    pro_completions.PRO_STATES.clear()
-
-
-def teardown_function():
-    pro_completions.PRO_STATES.clear()
-
-
 def test_pro_queue_full_returns_429(monkeypatch):
-    monkeypatch.setattr(pro_completions, "PRO_QUEUE_MAX_SIZE", 0)
-    client = TestClient(main.app)
-    response = client.post("/v1/pro/completions", headers=HEADERS, json=_payload())
+    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
-    assert response.json()["error"] == "PRO queue is full"
 
 
 def test_pro_status_missing_returns_404():
-    client = TestClient(main.app)
-    response = client.get("/v1/pro/completions/status/missing", headers=HEADERS)
+    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 = pro_completions._build_pro_prompts(
+    system_prompt, user_prompt = prompt.build_pro_completion_prompts(
         prefix="欢迎使用 LLM-IN-TEXT\n\n即时可用的 LLM 系统",
         suffix="",
         language_id="markdown",
         instruction="",
-        pro_thinking="high",
+        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
-    assert "long paragraphs or section-level output are allowed" in combined
-    assert "highest priority" in combined
-    assert "never copy tags to output" in combined
-    assert "write only the markdown that belongs at the cursor" not in combined
-    assert "continue the markdown naturally" in combined
 
 
-def test_pro_cancel_waits_for_stream_cleanup(monkeypatch):
-    started = threading.Event()
-    cleaned = threading.Event()
-
+def test_pro_stream_returns_standard_events(monkeypatch):
     async def fake_stream_events(*args, **kwargs):
-        started.set()
-        try:
-            yield "thinking", ""
-            while True:
-                await asyncio.sleep(0.05)
-        finally:
-            cleaned.set()
-
-    monkeypatch.setattr(pro_completions, "stream_ollama_events", fake_stream_events)
-    request_id = "pro-cancel-cleanup"
-    headers = {**HEADERS, "X-Request-Id": request_id}
-    response_box = {}
+        yield "thinking", ""
+        yield "content", "深度"
+        yield "content", "回答"
 
+    monkeypatch.setattr(job_handlers, "stream_ollama_events", fake_stream_events)
     with TestClient(main.app) as client:
-        def send_stream():
-            with client.stream("POST", "/v1/pro/completions", headers=headers, json=_payload()) as response:
-                response_box["status_code"] = response.status_code
-                response_box["body"] = "".join(response.iter_text())
+        with client.stream("POST", "/v1/pro/completions", headers=HEADERS, json=_payload()) as resp:
+            assert resp.status_code == 200
+            body = "".join(resp.iter_text())
 
-        stream_thread = threading.Thread(target=send_stream, daemon=True)
-        stream_thread.start()
-
-        assert started.wait(timeout=2.0)
-        cancel_response = client.post(
-            "/v1/pro/completions/cancel",
-            headers=HEADERS,
-            json={"request_id": request_id, "reason": "test"},
-        )
-
-        assert cancel_response.status_code == 200
-        assert cancel_response.json() == {"cancelled": True, "status": "ok"}
-        assert cleaned.wait(timeout=2.0)
-
-        stream_thread.join(timeout=5.0)
-        assert not stream_thread.is_alive()
+    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