llm-in-text/backend/tests/test_pro_completions.py

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,
        "user_preferences": {
            "language": "zh",
            "country": "CN",
            "timezone": "Asia/Shanghai",
        },
    }


def test_pro_queue_full_returns_429(monkeypatch):
    async def fake_queue_job(*args, **kwargs):
        raise job_system.QueueFullError("pro_completion", 8)

    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_prompt_accepts_serialized_preferences():
    _, user_prompt = prompt.build_pro_completion_prompts(
        prefix="Before",
        suffix="After",
        language_id="markdown",
        instruction="expand",
        preferences={
            "language": "zh",
            "country": "CN",
            "timezone": "Asia/Shanghai",
        },
    )

    assert "Preferred language: zh" in user_prompt
    assert "Preferred country: CN" in user_prompt
    assert "Preferred timezone: Asia/Shanghai" in user_prompt


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