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

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
import risk_control  # type: ignore
import session_store  # type: ignore
import audit_store  # 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()
    risk_control.reset_risk_controller()
    session_store.reset_session_store()
    audit_store.reset_audit_store()
    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"}