“ydy0615”
96 lines
2.7 KiB
Python
96 lines
2.7 KiB
Python
import asyncio
|
|
import importlib
|
|
import os
|
|
import sys
|
|
import threading
|
|
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
|
|
|
|
main = importlib.import_module("main")
|
|
|
|
API_KEY_HEADERS = {"X-API-Key": "your-secret-key-here"}
|
|
|
|
|
|
def setup_function():
|
|
job_system.reset_job_manager()
|
|
main._handlers_registered = False
|
|
|
|
|
|
def _completion_payload():
|
|
return {
|
|
"prefix": "hello",
|
|
"suffix": "",
|
|
"languageId": "markdown",
|
|
"model_thinking": "low",
|
|
"privacy_mode": True,
|
|
}
|
|
|
|
|
|
def test_cancel_endpoint_cancels_running_task(monkeypatch):
|
|
started = threading.Event()
|
|
cancelled = threading.Event()
|
|
|
|
async def fake_call_ollama(*args, **kwargs):
|
|
started.set()
|
|
try:
|
|
while True:
|
|
await asyncio.sleep(0.05)
|
|
except asyncio.CancelledError:
|
|
cancelled.set()
|
|
raise
|
|
|
|
monkeypatch.setattr(job_handlers, "call_ollama", fake_call_ollama)
|
|
request_id = "req-cancel-1"
|
|
|
|
with TestClient(main.app) as client:
|
|
response_box = {}
|
|
|
|
def send_completion():
|
|
with client.stream(
|
|
"POST",
|
|
"/v1/completions",
|
|
headers={**API_KEY_HEADERS, "X-Request-Id": request_id},
|
|
json=_completion_payload(),
|
|
) as response:
|
|
response_box["status_code"] = response.status_code
|
|
response_box["body"] = "".join(response.iter_text())
|
|
|
|
completion_thread = threading.Thread(target=send_completion, daemon=True)
|
|
completion_thread.start()
|
|
|
|
assert started.wait(timeout=2.0)
|
|
|
|
cancel_response = client.post(
|
|
"/v1/completions/cancel",
|
|
headers=API_KEY_HEADERS,
|
|
json={"request_id": request_id, "reason": "superseded"},
|
|
)
|
|
assert cancel_response.status_code == 200
|
|
assert cancel_response.json() == {"cancelled": True, "status": "ok"}
|
|
|
|
completion_thread.join(timeout=5.0)
|
|
assert not completion_thread.is_alive()
|
|
assert cancelled.wait(timeout=2.0)
|
|
assert "event: cancelled" in response_box["body"]
|
|
|
|
|
|
def test_cancel_not_found():
|
|
with TestClient(main.app) as client:
|
|
response = client.post(
|
|
"/v1/completions/cancel",
|
|
headers=API_KEY_HEADERS,
|
|
json={"request_id": "missing", "reason": "abort"},
|
|
)
|
|
assert response.status_code == 200
|
|
assert response.json() == {"cancelled": False, "status": "not_found"}
|