Migrate backend jobs to Redis Streams
This commit is contained in:
“ydy0615”
@@ -0,0 +1,79 @@
|
||||
import importlib
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
os.environ["JOB_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
|
||||
|
||||
main = importlib.import_module("main")
|
||||
|
||||
API_KEY = main.API_KEY
|
||||
HEADERS = {"X-API-Key": API_KEY}
|
||||
|
||||
|
||||
def setup_function():
|
||||
job_system.reset_job_manager()
|
||||
main._handlers_registered = False
|
||||
|
||||
|
||||
def _submit(client, content="test document", doc_type="txt", headers=None):
|
||||
return client.post("/v1/compress/submit", headers=headers if headers is not None else HEADERS, json={
|
||||
"content": content,
|
||||
"docType": doc_type,
|
||||
})
|
||||
|
||||
|
||||
def _status(client, task_id, headers=None):
|
||||
return client.get(f"/v1/compress/status?task_id={task_id}", headers=headers if headers is not None else HEADERS)
|
||||
|
||||
|
||||
def test_submit_empty_content_returns_400():
|
||||
with TestClient(main.app) as client:
|
||||
resp = _submit(client, "")
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
def test_submit_too_long_returns_400(monkeypatch):
|
||||
monkeypatch.setattr(main, "DOC_COMPRESS_CONTEXT_LIMIT", 10)
|
||||
with TestClient(main.app) as client:
|
||||
resp = _submit(client, "a" * 100)
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
def test_submit_success_returns_task_id():
|
||||
with TestClient(main.app) as client:
|
||||
resp = _submit(client, "hello world")
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert "task_id" in data
|
||||
assert data["status"] == "queued"
|
||||
|
||||
|
||||
def test_status_not_found_returns_404():
|
||||
with TestClient(main.app) as client:
|
||||
resp = _status(client, "nonexistent-id")
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
def test_status_completed(monkeypatch):
|
||||
async def fake_call_ollama(prompt, system_prompt=None, **kwargs): # noqa: ARG001
|
||||
return {"content": f"[compressed] {prompt[:20]}"}
|
||||
|
||||
monkeypatch.setattr(job_handlers, "call_ollama", fake_call_ollama)
|
||||
with TestClient(main.app) as client:
|
||||
resp = _submit(client, "important document text")
|
||||
task_id = resp.json()["task_id"]
|
||||
status_resp = _status(client, task_id)
|
||||
data = status_resp.json()
|
||||
assert status_resp.status_code == 200
|
||||
assert data["status"] in {"queued", "processing", "completed"}
|
||||
@@ -1,26 +1,31 @@
|
||||
import asyncio
|
||||
import importlib
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
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))
|
||||
|
||||
try:
|
||||
main = importlib.import_module("main")
|
||||
except ModuleNotFoundError:
|
||||
pytest.skip("main module dependencies are not available", allow_module_level=True)
|
||||
import job_handlers # type: ignore
|
||||
import job_system # type: ignore
|
||||
|
||||
main = importlib.import_module("main")
|
||||
|
||||
API_KEY_HEADERS = {"X-API-Key": "your-secret-key-here"}
|
||||
|
||||
|
||||
def setup_function():
|
||||
job_system.reset_job_manager()
|
||||
main._handlers_registered = False
|
||||
|
||||
|
||||
def _completion_payload():
|
||||
return {
|
||||
"prefix": "hello",
|
||||
@@ -32,7 +37,6 @@ def _completion_payload():
|
||||
|
||||
|
||||
def test_cancel_endpoint_cancels_running_task(monkeypatch):
|
||||
main.ACTIVE_COMPLETIONS.clear()
|
||||
started = threading.Event()
|
||||
cancelled = threading.Event()
|
||||
|
||||
@@ -45,21 +49,21 @@ def test_cancel_endpoint_cancels_running_task(monkeypatch):
|
||||
cancelled.set()
|
||||
raise
|
||||
|
||||
monkeypatch.setattr(main, "call_ollama", fake_call_ollama)
|
||||
monkeypatch.setattr(main, "build_completion_prompts", lambda *a, **k: ("system", "user"))
|
||||
monkeypatch.setattr(main, "prepare_prompt_context", lambda *a, **k: ("prefix", "suffix"))
|
||||
monkeypatch.setattr(job_handlers, "call_ollama", fake_call_ollama)
|
||||
request_id = "req-cancel-1"
|
||||
|
||||
with TestClient(main.app) as client:
|
||||
request_id = "req-cancel-1"
|
||||
completion_headers = {**API_KEY_HEADERS, "X-Request-Id": request_id}
|
||||
response_box = {}
|
||||
|
||||
def send_completion():
|
||||
response_box["response"] = client.post(
|
||||
with client.stream(
|
||||
"POST",
|
||||
"/v1/completions",
|
||||
headers=completion_headers,
|
||||
headers={**API_KEY_HEADERS, "X-Request-Id": request_id},
|
||||
json=_completion_payload(),
|
||||
)
|
||||
) as response:
|
||||
response_box["status_code"] = response.status_code
|
||||
response_box["body"] = "".join(response.iter_text())
|
||||
|
||||
completion_thread = threading.Thread(target=send_completion, daemon=True)
|
||||
completion_thread.start()
|
||||
@@ -77,16 +81,10 @@ def test_cancel_endpoint_cancels_running_task(monkeypatch):
|
||||
completion_thread.join(timeout=5.0)
|
||||
assert not completion_thread.is_alive()
|
||||
assert cancelled.wait(timeout=2.0)
|
||||
|
||||
completion_response = response_box["response"]
|
||||
# 499 = client disconnected (TestClient timeout during cancel)
|
||||
assert completion_response.status_code in (200, 499)
|
||||
if completion_response.status_code == 200:
|
||||
assert completion_response.json()["cancelled"] is True
|
||||
assert "event: cancelled" in response_box["body"]
|
||||
|
||||
|
||||
def test_cancel_not_found():
|
||||
main.ACTIVE_COMPLETIONS.clear()
|
||||
with TestClient(main.app) as client:
|
||||
response = client.post(
|
||||
"/v1/completions/cancel",
|
||||
@@ -95,27 +93,3 @@ def test_cancel_not_found():
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"cancelled": False, "status": "not_found"}
|
||||
|
||||
|
||||
def test_completion_normal_flow(monkeypatch):
|
||||
main.ACTIVE_COMPLETIONS.clear()
|
||||
|
||||
async def fake_call_ollama(*args, **kwargs):
|
||||
return {"content": "completion text", "think": ""}
|
||||
|
||||
monkeypatch.setattr(main, "call_ollama", fake_call_ollama)
|
||||
monkeypatch.setattr(main, "build_completion_prompts", lambda *a, **k: ("system", "user"))
|
||||
monkeypatch.setattr(main, "prepare_prompt_context", lambda *a, **k: ("prefix", "suffix"))
|
||||
|
||||
with TestClient(main.app) as client:
|
||||
response = client.post(
|
||||
"/v1/completions",
|
||||
headers=API_KEY_HEADERS,
|
||||
json=_completion_payload(),
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["content"] == "completion text"
|
||||
assert data["request_id"] is not None
|
||||
assert main.ACTIVE_COMPLETIONS == {}
|
||||
|
||||
@@ -1,35 +1,29 @@
|
||||
import base64
|
||||
import importlib
|
||||
import os
|
||||
import sys
|
||||
import base64
|
||||
import types
|
||||
import pytest
|
||||
from unittest.mock import MagicMock
|
||||
from pathlib import Path
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
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)
|
||||
os.environ["JOB_BACKEND"] = "memory"
|
||||
|
||||
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
|
||||
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 main # type: ignore
|
||||
import pro_completions # type: ignore
|
||||
import job_handlers # type: ignore
|
||||
import job_system # type: ignore
|
||||
|
||||
API_KEY = main.API_KEY
|
||||
HEADERS = {"X-API-Key": API_KEY}
|
||||
main = importlib.import_module("main")
|
||||
|
||||
HEADERS = {"X-API-Key": main.API_KEY}
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_active_completions():
|
||||
main.ACTIVE_COMPLETIONS.clear()
|
||||
pro_completions.PRO_STATES.clear()
|
||||
yield
|
||||
main.ACTIVE_COMPLETIONS.clear()
|
||||
pro_completions.PRO_STATES.clear()
|
||||
def setup_function():
|
||||
job_system.reset_job_manager()
|
||||
main._handlers_registered = False
|
||||
|
||||
|
||||
class DummyRequest:
|
||||
@@ -51,64 +45,12 @@ def test_preview_long_text_truncated():
|
||||
assert main._preview(long_text) == long_text[:80] + "..."
|
||||
|
||||
|
||||
def test_preview_none_input():
|
||||
assert main._preview(None) == ""
|
||||
|
||||
|
||||
def test_preview_newlines_replaced():
|
||||
assert main._preview("line1\nline2") == "line1\\nline2"
|
||||
|
||||
|
||||
def test_sanitize_markdown_strips_image_markdown():
|
||||
assert "" not in main._sanitize_converted_markdown(
|
||||
"text with image  end"
|
||||
)
|
||||
|
||||
|
||||
def test_sanitize_markdown_strips_img_tag():
|
||||
assert "<img" not in main._sanitize_converted_markdown("<img src='x.png'/>")
|
||||
|
||||
|
||||
def test_sanitize_markdown_collapse_newlines():
|
||||
assert main._sanitize_converted_markdown("a\n\n\nb\n\n\n\nc") == "a\n\nb\n\nc"
|
||||
|
||||
|
||||
def test_sanitize_markdown_normalize_crlf():
|
||||
result = main._sanitize_converted_markdown("line1\r\nline2\r\n")
|
||||
assert "line1\nline2" in result
|
||||
assert "\r" not in result
|
||||
assert "" not in main._sanitize_converted_markdown("text ")
|
||||
|
||||
|
||||
def test_sanitize_inline_completion_strips_prefill():
|
||||
assert main.sanitize_inline_completion_content(
|
||||
"系统非常适合写作",
|
||||
prefill="系统",
|
||||
) == "非常适合写作"
|
||||
|
||||
|
||||
def test_sanitize_inline_completion_extracts_fim_middle():
|
||||
assert main.sanitize_inline_completion_content(
|
||||
"<|fim_middle|>系统非常适合写作<|end|>",
|
||||
prefill="系统",
|
||||
) == "非常适合写作"
|
||||
|
||||
|
||||
def test_sanitize_inline_completion_extracts_polluted_chat_output():
|
||||
polluted = (
|
||||
"on new line? Prefix ends with newline already. The suffix starts with no newline. "
|
||||
"We need to consider if output should end with newline? The suffix starts with no newline. "
|
||||
"So we output: \"让我们一起探索 AI 的无限可能。\""
|
||||
"<|end|><|start|>assistant<|channel|>final|fim_middle|>系统让我们一起探索 AI 的无限可能。"
|
||||
)
|
||||
assert main.sanitize_inline_completion_content(
|
||||
polluted,
|
||||
prefill="系统",
|
||||
) == "让我们一起探索 AI 的无限可能。"
|
||||
|
||||
|
||||
def test_get_client_ip_from_host():
|
||||
req = DummyRequest(host="1.2.3.4", headers={})
|
||||
assert main.get_client_ip(req) == "1.2.3.4"
|
||||
assert main.sanitize_inline_completion_content("系统非常适合写作", prefill="系统") == "非常适合写作"
|
||||
|
||||
|
||||
def test_get_client_ip_header_overrides_host():
|
||||
@@ -116,191 +58,64 @@ def test_get_client_ip_header_overrides_host():
|
||||
assert main.get_client_ip(req) == "5.6.7.8"
|
||||
|
||||
|
||||
def test_get_client_ip_when_client_missing():
|
||||
req = DummyRequest(host=None, headers={"X-Client-IP": "9.9.9.9"})
|
||||
req.client = None
|
||||
assert main.get_client_ip(req) == "9.9.9.9"
|
||||
|
||||
|
||||
def test_post_completions_wrong_api_key_returns_401():
|
||||
client = TestClient(main.app)
|
||||
resp = client.post("/v1/completions", json={
|
||||
"prefix": "hello", "suffix": "", "languageId": "markdown",
|
||||
"model_thinking": "low", "privacy_mode": True,
|
||||
})
|
||||
with TestClient(main.app) as client:
|
||||
resp = client.post("/v1/completions", json={
|
||||
"prefix": "hello", "suffix": "", "languageId": "markdown",
|
||||
"model_thinking": "low", "privacy_mode": True,
|
||||
})
|
||||
assert resp.status_code == 401
|
||||
|
||||
|
||||
def test_post_completions_privacy_mode(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
def test_post_completions_returns_sse_done(monkeypatch):
|
||||
async def fake_call(*args, **kwargs):
|
||||
captured["kwargs"] = kwargs
|
||||
return {"content": "done", "think": ""}
|
||||
monkeypatch.setattr(main, "call_ollama", fake_call)
|
||||
monkeypatch.setattr(main, "build_completion_prompts", lambda *a, **k: ("sys", "user", ""))
|
||||
monkeypatch.setattr(main, "prepare_prompt_context", lambda *a, **k: ("p", "s"))
|
||||
|
||||
client = TestClient(main.app)
|
||||
resp = client.post("/v1/completions", headers=HEADERS, json={
|
||||
"prefix": "hello", "suffix": "", "languageId": "markdown",
|
||||
"model_thinking": "low", "privacy_mode": True,
|
||||
})
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data.get("content") == "done"
|
||||
# enable_thinking removed in OpenAI-compatible rewrite
|
||||
assert captured["kwargs"]["thinking"] == "low"
|
||||
|
||||
|
||||
def test_old_post_pro_stream_returns_404():
|
||||
client = TestClient(main.app)
|
||||
resp = client.post("/v1/pro/completions/stream", headers=HEADERS, json={
|
||||
"prefix": "hello",
|
||||
"suffix": "",
|
||||
"languageId": "markdown",
|
||||
"model_thinking": "high",
|
||||
"privacy_mode": True,
|
||||
})
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
def test_post_pro_completion_returns_sse_and_status(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
async def fake_stream_events(*args, **kwargs):
|
||||
captured["kwargs"] = kwargs
|
||||
yield "thinking", ""
|
||||
yield "content", "深度"
|
||||
yield "content", "回答"
|
||||
|
||||
monkeypatch.setattr(pro_completions, "stream_ollama_events", fake_stream_events)
|
||||
client = TestClient(main.app)
|
||||
with client.stream("POST", "/v1/pro/completions", headers=HEADERS, json={
|
||||
"prefix": "hello",
|
||||
"suffix": "",
|
||||
"languageId": "markdown",
|
||||
"instruction": "expand",
|
||||
"pro_thinking": "high",
|
||||
"privacy_mode": True,
|
||||
}) as resp:
|
||||
assert resp.status_code == 200
|
||||
body = "".join(resp.iter_text())
|
||||
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: thinking" in body
|
||||
assert "event: chunk" in body
|
||||
assert "event: result" in body
|
||||
assert "event: done" in body
|
||||
assert "深度" in body
|
||||
assert "回答" in body
|
||||
assert captured["kwargs"]["use_pro_model"] is True
|
||||
assert captured["kwargs"]["thinking"] == "high"
|
||||
|
||||
request_id = next(iter(pro_completions.PRO_STATES))
|
||||
status_resp = client.get(f"/v1/pro/completions/status/{request_id}", headers=HEADERS)
|
||||
assert status_resp.status_code == 200
|
||||
assert status_resp.json()["status"] == "done"
|
||||
assert main.ACTIVE_COMPLETIONS == {}
|
||||
|
||||
|
||||
def test_post_ocr_mocked(monkeypatch):
|
||||
async def fake_ocr(*args, **kwargs):
|
||||
return "OCR result text"
|
||||
monkeypatch.setattr(main, "call_vlm_ocr", fake_ocr)
|
||||
|
||||
client = TestClient(main.app)
|
||||
monkeypatch.setattr(job_handlers, "call_vlm_ocr", fake_ocr)
|
||||
img_b64 = base64.b64encode(b"pretend image data").decode()
|
||||
resp = client.post("/v1/ocr", headers=HEADERS, json={
|
||||
"image": img_b64, "filename": "test.jpg", "language": "auto",
|
||||
})
|
||||
assert resp.status_code == 200
|
||||
j = resp.json()
|
||||
assert j["text"] == "OCR result text"
|
||||
assert j["filename"] == "test.jpg"
|
||||
|
||||
|
||||
def test_post_ocr_invalid_base64_returns_500():
|
||||
client = TestClient(main.app)
|
||||
resp = client.post("/v1/ocr", headers=HEADERS, json={
|
||||
"image": "not-base64!!!", "filename": "test.jpg",
|
||||
})
|
||||
assert resp.status_code == 500
|
||||
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_convert_txt_returns_markdown():
|
||||
client = TestClient(main.app)
|
||||
content = base64.b64encode(b"hello world").decode()
|
||||
resp = client.post("/v1/convert", headers=HEADERS, json={
|
||||
"file": content, "filename": "sample.txt",
|
||||
})
|
||||
assert resp.status_code == 200
|
||||
j = resp.json()
|
||||
assert j["markdown"] == "hello world"
|
||||
assert j["filename"] == "sample.txt"
|
||||
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():
|
||||
client = TestClient(main.app)
|
||||
content = base64.b64encode(b"data").decode()
|
||||
resp = client.post("/v1/convert", headers=HEADERS, json={
|
||||
"file": content, "filename": "sample.xlsx",
|
||||
})
|
||||
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_post_convert_docx_with_mocked_markitdown(monkeypatch):
|
||||
class FakeResult:
|
||||
text_content = "markdown from docx"
|
||||
class FakeMD:
|
||||
def convert(self, path):
|
||||
return FakeResult()
|
||||
monkeypatch.setattr(main, "_get_markitdown", lambda: FakeMD())
|
||||
|
||||
client = TestClient(main.app)
|
||||
content = base64.b64encode(b"docx content").decode()
|
||||
resp = client.post("/v1/convert", headers=HEADERS, json={
|
||||
"file": content, "filename": "sample.docx",
|
||||
})
|
||||
assert resp.status_code == 200
|
||||
j = resp.json()
|
||||
assert j["markdown"] == "markdown from docx"
|
||||
|
||||
|
||||
def test_post_cancel_non_existent_returns_not_found():
|
||||
client = TestClient(main.app)
|
||||
resp = client.post("/v1/completions/cancel", headers=HEADERS, json={
|
||||
"request_id": "non-existent", "reason": "abort",
|
||||
})
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["cancelled"] is False
|
||||
assert data["status"] == "not_found"
|
||||
|
||||
|
||||
def test_post_cancel_wrong_api_key_returns_401():
|
||||
client = TestClient(main.app)
|
||||
resp = client.post("/v1/completions/cancel", json={
|
||||
"request_id": "id", "reason": "abort",
|
||||
})
|
||||
assert resp.status_code == 401
|
||||
|
||||
|
||||
def test_post_cancel_already_done(monkeypatch):
|
||||
main.ACTIVE_COMPLETIONS.clear()
|
||||
# Create a mock task that appears done
|
||||
mock_task = MagicMock()
|
||||
mock_task.done.return_value = True
|
||||
mock_task.cancel = MagicMock()
|
||||
main.ACTIVE_COMPLETIONS["done-id"] = mock_task
|
||||
|
||||
client = TestClient(main.app)
|
||||
resp = client.post("/v1/completions/cancel", headers=HEADERS, json={
|
||||
"request_id": "done-id", "reason": "abort",
|
||||
})
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["cancelled"] is False
|
||||
assert data["status"] == "already_done"
|
||||
main.ACTIVE_COMPLETIONS.clear()
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user