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

import base64
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")

HEADERS = {"X-API-Key": main.API_KEY}


def setup_function():
    job_system.reset_job_manager()
    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_wrong_api_key_returns_401():
    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_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_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_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"]