80 lines
2.3 KiB
Python
80 lines
2.3 KiB
Python
|
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"}
|