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

import base64
import asyncio
import base64
import importlib
import os
import sys
from pathlib import Path
from types import SimpleNamespace

from fastapi.testclient import TestClient

os.environ["JOB_BACKEND"] = "memory"
os.environ["DOCS_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
import docs_store  # type: ignore
import risk_control  # type: ignore
import session_store  # type: ignore
import audit_store  # type: ignore

main = importlib.import_module("main")

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


def setup_function():
    job_system.reset_job_manager()
    docs_store.reset_document_store()
    risk_control.reset_risk_controller()
    session_store.reset_session_store()
    audit_store.reset_audit_store()
    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_without_api_key_uses_anonymous_session(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", json={
            "prefix": "hello", "suffix": "", "languageId": "markdown",
            "model_thinking": "low", "privacy_mode": True,
        }) as resp:
            assert resp.status_code == 200
            assert main.config.session_cookie_name in resp.cookies


def test_post_completions_invalid_api_key_returns_403():
    with TestClient(main.app) as client:
        resp = client.post(
            "/v1/completions",
            headers={"X-API-Key": "invalid-key"},
            json={
                "prefix": "hello", "suffix": "", "languageId": "markdown",
                "model_thinking": "low", "privacy_mode": True,
            },
        )
    assert resp.status_code == 403


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_stream_job_emits_keepalive_during_idle(monkeypatch):
    async def fake_queue_job(*_args, **_kwargs):
        return "job-keepalive"

    class FakeManager:
        def register_handler(self, *_args, **_kwargs):
            return None

        async def stream_events(self, job_id):
            yield {"event": "queued", "job_id": job_id}
            yield {"event": "started", "job_id": job_id}
            await asyncio.sleep(0.03)
            yield {"event": "done", "job_id": job_id, "result": {"content": "ok"}}

    monkeypatch.setattr(main, "STREAM_HEARTBEAT_SECONDS", 0.01)
    monkeypatch.setattr(main, "_queue_job", fake_queue_job)
    monkeypatch.setattr(main, "get_job_manager", lambda: FakeManager())

    with TestClient(main.app) as client:
        with client.stream("POST", "/v1/pro/completions", headers=HEADERS, json={
            "prefix": "hello", "suffix": "", "languageId": "markdown",
            "instruction": "expand", "pro_thinking": "medium", "privacy_mode": True,
        }) as resp:
            assert resp.status_code == 200
            body = "".join(resp.iter_text())

    assert ": keepalive" 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_video_ocr_merges_ocr_and_asr(monkeypatch):
    async def fake_ocr(*args, **kwargs):
        return "画面文字"

    async def fake_asr(*args, **kwargs):
        return SimpleNamespace(text="音频转写")

    monkeypatch.setattr(job_handlers, "call_vlm_ocr", fake_ocr)
    monkeypatch.setattr(job_handlers, "generate_asr_response", fake_asr)
    monkeypatch.setattr(job_handlers, "extract_audio_wav_bytes", lambda _path: b"fake wav")

    video_b64 = base64.b64encode(b"pretend video data").decode()
    with TestClient(main.app) as client:
        with client.stream("POST", "/v1/ocr", headers=HEADERS, json={
            "image": video_b64,
            "filename": "sample.mp4",
            "language": "auto",
            "media_type": "video",
            "mime_type": "video/mp4",
        }) as resp:
            assert resp.status_code == 200
            body = "".join(resp.iter_text())

    assert "视频画面 OCR" in body
    assert "视频音频 ASR" in body
    assert "画面文字" in body
    assert "音频转写" 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"]


def test_post_convert_rejects_mismatched_content_suffix():
    content = base64.b64encode(b"%PDF-1.4\n%%EOF").decode()
    with TestClient(main.app) as client:
        resp = client.post("/v1/convert", headers=HEADERS, json={
            "file": content, "filename": "sample.txt",
        })
    assert resp.status_code == 500
    assert "仅支持" in resp.json()["error"]


def test_docs_nodes_crud_round_trip():
    with TestClient(main.app) as client:
        folder_resp = client.post("/v1/docs/folders", headers=HEADERS, json={
            "name": "项目资料",
            "parentId": None,
        })
        assert folder_resp.status_code == 200
        folder = folder_resp.json()["node"]

        file_resp = client.post("/v1/docs/files/text", headers=HEADERS, json={
            "name": "notes.md",
            "parentId": folder["id"],
            "content": "# hello",
        })
        assert file_resp.status_code == 200
        file_node = file_resp.json()["node"]
        assert file_node["previewText"] == "# hello"

        list_resp = client.get("/v1/docs/nodes", headers=HEADERS)
        assert list_resp.status_code == 200
        nodes = list_resp.json()["nodes"]
        assert len(nodes) == 2

        rename_resp = client.patch(f"/v1/docs/nodes/{file_node['id']}", headers=HEADERS, json={
            "name": "renamed.md",
        })
        assert rename_resp.status_code == 200
        assert rename_resp.json()["node"]["name"] == "renamed.md"

        blob_resp = client.get(f"/v1/docs/files/{file_node['id']}/blob", headers=HEADERS)
        assert blob_resp.status_code == 200
        assert blob_resp.content == b"# hello"

        delete_resp = client.delete(f"/v1/docs/nodes/{folder['id']}", headers=HEADERS)
        assert delete_resp.status_code == 200

        final_list = client.get("/v1/docs/nodes", headers=HEADERS)
        assert final_list.status_code == 200
        assert final_list.json()["nodes"] == []


def test_docs_file_upload_and_blob_replace():
    with TestClient(main.app) as client:
        upload_resp = client.post(
            "/v1/docs/files/upload",
            headers=HEADERS,
            files={"file": ("image.png", b"png-bytes", "image/png")},
            data={"parent_id": ""},
        )
        assert upload_resp.status_code == 200
        node = upload_resp.json()["node"]
        assert node["storageKind"] == "blob"

        replace_resp = client.put(
            f"/v1/docs/files/{node['id']}/blob",
            headers=HEADERS,
            files={"file": ("photo.jpg", b"jpeg-bytes", "image/jpeg")},
        )
        assert replace_resp.status_code == 200
        assert replace_resp.json()["node"]["name"] == "photo.jpg"

        blob_resp = client.get(f"/v1/docs/files/{node['id']}/blob", headers=HEADERS)
        assert blob_resp.status_code == 200
        assert blob_resp.content == b"jpeg-bytes"