thePlug/backend/tests/services/test_health_monitor.py

import asyncio
from datetime import UTC, datetime

import pytest

from app.models.health import HealthSnapshot
from app.models.proxy import ProxyDocument
from app.models.service import ServiceEndpointDocument
from app.services.mitmproxy.health_monitor import HealthMonitor


def make_proxy(**updates):
    data = {
        "proxy_port": 8080,
        "name": "svc",
        "code": "",
        "status": "running",
        "service": ServiceEndpointDocument(port=9000),
    }
    data.update(updates)
    return ProxyDocument(**data)


def make_snapshot(proxy_port: int, *, overall: str = "healthy") -> HealthSnapshot:
    return HealthSnapshot(
        proxy_port=proxy_port,
        timestamp=datetime.now(UTC),
        enabled=True,
        overall=overall,
        checks=[],
    )


def test_run_cycle_records_and_publishes_running_healthchecks(monkeypatch):
    monitor = HealthMonitor()
    proxy = make_proxy()
    captured: dict[str, object] = {}
    snapshot = make_snapshot(8080)

    async def fake_list_proxies():
        return [proxy]

    async def fake_check_proxy(document):
        captured["checked"] = document.proxy_port
        return snapshot

    async def fake_record_health_sample(document):
        captured["recorded"] = document
        return document

    async def fake_publish(proxy_port, document):
        captured["published"] = (proxy_port, document)

    monkeypatch.setattr(
        "app.services.mitmproxy.health_monitor.list_proxies",
        fake_list_proxies,
    )
    monkeypatch.setattr(
        "app.services.mitmproxy.health_monitor.check_proxy",
        fake_check_proxy,
    )
    monkeypatch.setattr(
        "app.services.mitmproxy.health_monitor.record_health_sample",
        fake_record_health_sample,
    )
    monkeypatch.setattr(monitor, "publish", fake_publish)

    asyncio.run(monitor.run_cycle())

    assert captured["checked"] == 8080
    assert captured["recorded"] == snapshot
    assert captured["published"] == (8080, snapshot)


def test_run_cycle_skips_stopped_or_disabled_proxies(monkeypatch):
    monitor = HealthMonitor()
    running_disabled = make_proxy(healthcheck={"enabled": False})
    stopped = make_proxy(proxy_port=8081, status="stopped")
    called = False

    async def fake_list_proxies():
        return [running_disabled, stopped]

    async def fake_check_proxy(_document):
        nonlocal called
        called = True
        return make_snapshot(8080)

    monkeypatch.setattr(
        "app.services.mitmproxy.health_monitor.list_proxies",
        fake_list_proxies,
    )
    monkeypatch.setattr(
        "app.services.mitmproxy.health_monitor.check_proxy",
        fake_check_proxy,
    )

    asyncio.run(monitor.run_cycle())

    assert called is False


def test_run_cycle_checks_proxies_in_parallel(monkeypatch):
    monitor = HealthMonitor()
    proxies = [
        make_proxy(proxy_port=8080),
        make_proxy(proxy_port=8081),
        make_proxy(proxy_port=8082),
    ]
    started: list[int] = []
    release = asyncio.Event()

    async def fake_list_proxies():
        return proxies

    async def fake_check_proxy(document):
        started.append(document.proxy_port)
        if len(started) == len(proxies):
            release.set()
        await release.wait()
        return make_snapshot(document.proxy_port)

    async def fake_record_health_sample(document):
        return document

    async def fake_publish(_proxy_port, _document):
        return None

    monkeypatch.setattr(
        "app.services.mitmproxy.health_monitor.list_proxies",
        fake_list_proxies,
    )
    monkeypatch.setattr(
        "app.services.mitmproxy.health_monitor.check_proxy",
        fake_check_proxy,
    )
    monkeypatch.setattr(
        "app.services.mitmproxy.health_monitor.record_health_sample",
        fake_record_health_sample,
    )
    monkeypatch.setattr(monitor, "publish", fake_publish)

    asyncio.run(asyncio.wait_for(monitor.run_cycle(), timeout=1))

    assert started == [8080, 8081, 8082]


def test_run_cycle_continues_when_one_proxy_check_fails(monkeypatch):
    monitor = HealthMonitor()
    proxies = [
        make_proxy(proxy_port=8080),
        make_proxy(proxy_port=8081),
    ]
    captured: list[int] = []

    async def fake_list_proxies():
        return proxies

    async def fake_check_proxy(document):
        if document.proxy_port == 8080:
            raise RuntimeError("boom")
        return make_snapshot(document.proxy_port)

    async def fake_record_health_sample(document):
        captured.append(document.proxy_port)
        return document

    async def fake_publish(proxy_port, _document):
        captured.append(proxy_port + 10_000)

    monkeypatch.setattr(
        "app.services.mitmproxy.health_monitor.list_proxies",
        fake_list_proxies,
    )
    monkeypatch.setattr(
        "app.services.mitmproxy.health_monitor.check_proxy",
        fake_check_proxy,
    )
    monkeypatch.setattr(
        "app.services.mitmproxy.health_monitor.record_health_sample",
        fake_record_health_sample,
    )
    monkeypatch.setattr(monitor, "publish", fake_publish)

    asyncio.run(monitor.run_cycle())

    assert captured == [8081, 18081]


def test_run_retries_after_cycle_failure(monkeypatch):
    monitor = HealthMonitor()
    calls = {"run_cycle": 0}
    sleep_calls: list[float] = []

    async def fake_run_cycle():
        calls["run_cycle"] += 1
        if calls["run_cycle"] == 1:
            raise RuntimeError("temporary failure")
        raise asyncio.CancelledError

    async def fake_sleep(seconds):
        sleep_calls.append(seconds)
        return None

    monkeypatch.setattr(monitor, "run_cycle", fake_run_cycle)
    monkeypatch.setattr(
        "app.services.mitmproxy.health_monitor.asyncio.sleep",
        fake_sleep,
    )

    try:
        asyncio.run(monitor.run())
    except asyncio.CancelledError:
        pass

    assert calls["run_cycle"] == 2
    assert len(sleep_calls) == 1
    assert sleep_calls[0] == pytest.approx(1.0, abs=0.01)