“ydy0615”
122 lines
3.9 KiB
Python
122 lines
3.9 KiB
Python
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 "" not in main._sanitize_converted_markdown("text ")
|
|
|
|
|
|
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"]
|