From f63e33eef2a43932bf2297282cceb40de832aa77 Mon Sep 17 00:00:00 2001
From: Roman Meshko <roman.meshko@gmail.com>
Date: Fri, 8 May 2026 17:25:00 +0300
Subject: [PATCH] Fixed isolate Jetson manual e2e and stabilize SSE warmup

---
 e2e/conftest.py             | 50 ++++++++++++++++++++++++++++++++++---
 e2e/docker-compose.test.yml |  3 +++
 e2e/tests/test_async_sse.py |  4 +--
 src/main.py                 |  1 +
 4 files changed, 52 insertions(+), 6 deletions(-)

diff --git a/e2e/conftest.py b/e2e/conftest.py
index 6096d52..bdaf093 100644
--- a/e2e/conftest.py
+++ b/e2e/conftest.py
@@ -133,10 +133,11 @@ def image_detect(http_client, auth_headers):
             headers=headers,
             timeout=timeout,
         )
-        done.wait(timeout=timeout)
+        completed = done.wait(timeout=timeout)
         elapsed_ms = (time.perf_counter() - t0) * 1000.0
 
         assert r.status_code == 202, f"Expected 202, got {r.status_code}: {r.text}"
+        assert completed, f"Timed out waiting for terminal SSE event on channel {cid}"
         assert not errors, f"SSE errors: {errors}"
 
         th.join(timeout=1)
@@ -214,6 +215,27 @@ def _read_media(name: str) -> bytes:
     return p.read_bytes()
 
 
+def _health(http_client):
+    r = http_client.get("/health")
+    r.raise_for_status()
+    return r.json()
+
+
+def _health_ai_active(data: dict) -> bool:
+    return data.get("aiAvailability") not in ("None", "Downloading", "Error")
+
+
+def _wait_for_ai_active(http_client, timeout: float = 30) -> dict | None:
+    deadline = time.time() + timeout
+    last = None
+    while time.time() < deadline:
+        last = _health(http_client)
+        if _health_ai_active(last):
+            return last
+        time.sleep(1)
+    return last
+
+
 @pytest.fixture(scope="session")
 def media_dir():
     return str(_media_dir())
@@ -280,11 +302,20 @@ def warm_engine(http_client, image_small, auth_headers):
     consecutive_errors = 0
 
     while time.time() < deadline:
+        try:
+            current_health = _health(http_client)
+            if _health_ai_active(current_health):
+                return
+        except OSError:
+            pass
+
         cid = str(uuid.uuid4())
         headers = {**auth_headers, "X-Channel-Id": cid}
         done = threading.Event()
+        terminal_status = None
 
         def _listen(cid=cid):
+            nonlocal terminal_status
             try:
                 with http_client.get(
                     f"/detect/events/{cid}",
@@ -297,7 +328,8 @@ def warm_engine(http_client, image_small, auth_headers):
                         if not ev.data or not str(ev.data).strip():
                             continue
                         data = json.loads(ev.data)
-                        if data.get("mediaStatus") == "AIProcessed":
+                        terminal_status = data.get("mediaStatus")
+                        if terminal_status in ("AIProcessed", "Error"):
                             break
             except Exception:
                 pass
@@ -316,9 +348,19 @@ def warm_engine(http_client, image_small, auth_headers):
             )
             last_status = r.status_code
             if r.status_code == 202:
-                done.wait(timeout=30)
+                completed = done.wait(timeout=30)
                 th.join(timeout=1)
-                return
+                if completed and terminal_status == "AIProcessed":
+                    if _health_ai_active(_wait_for_ai_active(http_client, timeout=30) or {}):
+                        return
+                if completed and terminal_status == "Error":
+                    consecutive_errors += 1
+                else:
+                    active_health = _wait_for_ai_active(http_client, timeout=10)
+                    if _health_ai_active(active_health or {}):
+                        return
+                time.sleep(2)
+                continue
             if r.status_code >= 500:
                 consecutive_errors += 1
                 if consecutive_errors >= 5:
diff --git a/e2e/docker-compose.test.yml b/e2e/docker-compose.test.yml
index 484d555..83da3fb 100644
--- a/e2e/docker-compose.test.yml
+++ b/e2e/docker-compose.test.yml
@@ -27,6 +27,7 @@ services:
       ANNOTATIONS_URL: http://mock-annotations:8081
       JWT_SECRET: test-secret-e2e-only
       CLASSES_JSON_PATH: /app/classes.json
+      LOG_DIR: /app/Logs
     volumes:
       - ./fixtures:/media
       - detections-logs:/app/Logs
@@ -50,6 +51,7 @@ services:
       ANNOTATIONS_URL: http://mock-annotations:8081
       JWT_SECRET: test-secret-e2e-only
       CLASSES_JSON_PATH: /app/classes.json
+      LOG_DIR: /app/Logs
     volumes:
       - ./fixtures:/media
       - detections-logs:/app/Logs
@@ -73,6 +75,7 @@ services:
       ANNOTATIONS_URL: http://mock-annotations:8081
       JWT_SECRET: test-secret-e2e-only
       CLASSES_JSON_PATH: /app/classes.json
+      LOG_DIR: /app/Logs
     volumes:
       - ./fixtures:/media:ro
       - detections-logs:/app/Logs
diff --git a/e2e/tests/test_async_sse.py b/e2e/tests/test_async_sse.py
index 5be048a..0feb038 100644
--- a/e2e/tests/test_async_sse.py
+++ b/e2e/tests/test_async_sse.py
@@ -29,11 +29,11 @@ def test_ft_p08_immediate_async_response(
     assert r.status_code == 202
 
 
-@pytest.mark.timeout(10)
+@pytest.mark.timeout(30)
 def test_ft_p09_sse_event_delivery(
     warm_engine, http_client, jwt_token
 ):
-    media_id = f"sse-{uuid.uuid4().hex}"
+    media_id = f"event-{uuid.uuid4().hex}"
     channel_id = str(uuid.uuid4())
     body = _ai_config_video()
     auth_header = {"Authorization": f"Bearer {jwt_token}"}
diff --git a/src/main.py b/src/main.py
index 5217704..c9a8d51 100644
--- a/src/main.py
+++ b/src/main.py
@@ -409,6 +409,7 @@ async def detect_events(channel_id: str, request: Request, after_ts: Optional[in
 
     async def event_generator():
         try:
+            yield ": connected\n\n"
             if after_ts is not None:
                 for ts_ms, data in list(_channel_buffers.get(channel_id, [])):
                     if ts_ms > after_ts: