detections/e2e/tests/test_health_engine.py

import json
import os
import threading
import time
import uuid

import pytest
import sseclient

_DETECT_TIMEOUT = 60
_ENGINE_WAIT_TIMEOUT = int(os.environ.get("E2E_ENGINE_WAIT_TIMEOUT", "900"))
_WAIT_FOR_ENGINE_ENABLED = os.environ.get("E2E_WAIT_FOR_ENGINE_ENABLED", "").lower() in (
    "1",
    "true",
    "yes",
    "on",
)


def _get_health(http_client):
    r = http_client.get("/health")
    r.raise_for_status()
    return r.json()


def _assert_active_ai(data):
    assert data["status"] == "healthy"
    assert data["aiAvailability"] not in ("None", "Downloading")


def _wait_for_engine_enabled(http_client):
    deadline = time.time() + _ENGINE_WAIT_TIMEOUT
    last = None
    while time.time() < deadline:
        last = _get_health(http_client)
        availability = last.get("aiAvailability")
        if availability == "Enabled":
            assert last.get("errorMessage") is None
            return last
        if availability == "Error":
            pytest.fail(f"engine conversion failed: {last.get('errorMessage')}")
        time.sleep(3)
    pytest.fail(
        f"engine did not become Enabled within {_ENGINE_WAIT_TIMEOUT}s "
        f"(last health: {last})"
    )


@pytest.mark.cpu
class TestHealthEngineStep01PreInit:
    def test_ft_p_01_pre_init_health(self, http_client):
        t0 = time.monotonic()
        data = _get_health(http_client)
        assert time.monotonic() - t0 < 2.0
        assert data["status"] == "healthy"
        if data["aiAvailability"] != "None":
            pytest.skip(
                f"engine already initialized (aiAvailability={data['aiAvailability']}); "
                "pre-init health check only applies to a cold service"
            )
        assert data["aiAvailability"] == "None", (
            f"engine already initialized (aiAvailability={data['aiAvailability']}); "
            "pre-init tests must run before any test that triggers warm_engine"
        )
        assert data.get("errorMessage") is None


@pytest.mark.cpu
@pytest.mark.slow
class TestHealthEngineStep02LazyInit:
    def test_ft_p_14_lazy_initialization(self, http_client, image_small, auth_headers):
        before = _get_health(http_client)
        if before["aiAvailability"] != "None":
            pytest.skip(
                f"engine already initialized (aiAvailability={before['aiAvailability']}); "
                "lazy-init check only applies to a cold service"
            )
        assert before["aiAvailability"] == "None", (
            f"engine already initialized (aiAvailability={before['aiAvailability']}); "
            "lazy-init test must run before any test that triggers warm_engine"
        )

        cid = str(uuid.uuid4())
        headers = {**auth_headers, "X-Channel-Id": cid}
        done = threading.Event()
        connected = threading.Event()

        def _listen():
            try:
                with http_client.get(
                    f"/detect/events/{cid}",
                    stream=True,
                    timeout=_DETECT_TIMEOUT + 2,
                    headers=auth_headers,
                ) as resp:
                    resp.raise_for_status()
                    connected.set()
                    for ev in sseclient.SSEClient(resp).events():
                        if not ev.data or not str(ev.data).strip():
                            continue
                        data = json.loads(ev.data)
                        if data.get("mediaStatus") in ("AIProcessed", "Error"):
                            break
            except Exception:
                pass
            finally:
                connected.set()
                done.set()

        th = threading.Thread(target=_listen, daemon=True)
        th.start()
        connected.wait(timeout=5)

        files = {"file": ("lazy.jpg", image_small, "image/jpeg")}
        r = http_client.post("/detect/image", files=files, headers=headers, timeout=_DETECT_TIMEOUT)
        assert r.status_code == 202
        done.wait(timeout=_DETECT_TIMEOUT)
        th.join(timeout=2)

        after = _get_health(http_client)
        _assert_active_ai(after)


@pytest.mark.cpu
@pytest.mark.slow
class TestHealthEngineStep03Warmed:
    @pytest.fixture(autouse=True)
    def _warm(self, warm_engine):
        pass

    @pytest.mark.timeout(_ENGINE_WAIT_TIMEOUT + 30)
    def test_ft_p_02_post_init_health(self, http_client):
        if _WAIT_FOR_ENGINE_ENABLED:
            data = _wait_for_engine_enabled(http_client)
        else:
            data = _get_health(http_client)
        _assert_active_ai(data)
        assert data.get("errorMessage") is None

    def test_ft_p_15_onnx_cpu_detect(self, http_client, image_small, auth_headers):
        cid = str(uuid.uuid4())
        headers = {**auth_headers, "X-Channel-Id": cid}
        all_detections = []
        done = threading.Event()
        connected = threading.Event()

        def _listen():
            try:
                with http_client.get(
                    f"/detect/events/{cid}",
                    stream=True,
                    timeout=_DETECT_TIMEOUT + 2,
                    headers=auth_headers,
                ) as resp:
                    resp.raise_for_status()
                    connected.set()
                    for ev in sseclient.SSEClient(resp).events():
                        if not ev.data or not str(ev.data).strip():
                            continue
                        data = json.loads(ev.data)
                        if data.get("mediaStatus") == "AIProcessing":
                            all_detections.extend(data.get("annotations", []))
                        if data.get("mediaStatus") in ("AIProcessed", "Error"):
                            break
            except Exception:
                pass
            finally:
                connected.set()
                done.set()

        th = threading.Thread(target=_listen, daemon=True)
        th.start()
        connected.wait(timeout=5)

        files = {"file": ("onnx.jpg", image_small, "image/jpeg")}
        r = http_client.post("/detect/image", files=files, headers=headers, timeout=_DETECT_TIMEOUT)
        r.raise_for_status()
        assert r.status_code == 202
        done.wait(timeout=_DETECT_TIMEOUT)
        th.join(timeout=2)

        if all_detections:
            d = all_detections[0]
            for k in (
                "centerX",
                "centerY",
                "width",
                "height",
                "classNum",
                "label",
                "confidence",
            ):
                assert k in d