Migrate backend jobs to Redis Streams

backend/tests/test_main_cancel.py

@@ -1,26 +1,31 @@
 import asyncio
 import importlib
+import os
 import sys
 import threading
 from pathlib import Path
 
-import pytest
 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))
 
-try:
-    main = importlib.import_module("main")
-except ModuleNotFoundError:
-    pytest.skip("main module dependencies are not available", allow_module_level=True)
+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",
@@ -32,7 +37,6 @@ def _completion_payload():
 
 
 def test_cancel_endpoint_cancels_running_task(monkeypatch):
-    main.ACTIVE_COMPLETIONS.clear()
     started = threading.Event()
     cancelled = threading.Event()
 
@@ -45,21 +49,21 @@ def test_cancel_endpoint_cancels_running_task(monkeypatch):
             cancelled.set()
             raise
 
-    monkeypatch.setattr(main, "call_ollama", fake_call_ollama)
-    monkeypatch.setattr(main, "build_completion_prompts", lambda *a, **k: ("system", "user"))
-    monkeypatch.setattr(main, "prepare_prompt_context", lambda *a, **k: ("prefix", "suffix"))
+    monkeypatch.setattr(job_handlers, "call_ollama", fake_call_ollama)
+    request_id = "req-cancel-1"
 
     with TestClient(main.app) as client:
-        request_id = "req-cancel-1"
-        completion_headers = {**API_KEY_HEADERS, "X-Request-Id": request_id}
         response_box = {}
 
         def send_completion():
-            response_box["response"] = client.post(
+            with client.stream(
+                "POST",
                 "/v1/completions",
-                headers=completion_headers,
+                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()
@@ -77,16 +81,10 @@ def test_cancel_endpoint_cancels_running_task(monkeypatch):
         completion_thread.join(timeout=5.0)
         assert not completion_thread.is_alive()
         assert cancelled.wait(timeout=2.0)
-
-        completion_response = response_box["response"]
-        # 499 = client disconnected (TestClient timeout during cancel)
-        assert completion_response.status_code in (200, 499)
-        if completion_response.status_code == 200:
-            assert completion_response.json()["cancelled"] is True
+        assert "event: cancelled" in response_box["body"]
 
 
 def test_cancel_not_found():
-    main.ACTIVE_COMPLETIONS.clear()
     with TestClient(main.app) as client:
         response = client.post(
             "/v1/completions/cancel",
@@ -95,27 +93,3 @@ def test_cancel_not_found():
         )
     assert response.status_code == 200
     assert response.json() == {"cancelled": False, "status": "not_found"}
-
-
-def test_completion_normal_flow(monkeypatch):
-    main.ACTIVE_COMPLETIONS.clear()
-
-    async def fake_call_ollama(*args, **kwargs):
-        return {"content": "completion text", "think": ""}
-
-    monkeypatch.setattr(main, "call_ollama", fake_call_ollama)
-    monkeypatch.setattr(main, "build_completion_prompts", lambda *a, **k: ("system", "user"))
-    monkeypatch.setattr(main, "prepare_prompt_context", lambda *a, **k: ("prefix", "suffix"))
-
-    with TestClient(main.app) as client:
-        response = client.post(
-            "/v1/completions",
-            headers=API_KEY_HEADERS,
-            json=_completion_payload(),
-        )
-
-    assert response.status_code == 200
-    data = response.json()
-    assert data["content"] == "completion text"
-    assert data["request_id"] is not None
-    assert main.ACTIVE_COMPLETIONS == {}