diff --git a/src/netcheck/orchestration.py b/src/netcheck/orchestration.py new file mode 100644 index 0000000..f4b3e8a --- /dev/null +++ b/src/netcheck/orchestration.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +import asyncio +import time +from datetime import datetime, timezone +from typing import Any, Awaitable + +import httpx + +from netcheck.models import ModuleResult, ProbeError + +_BLOCKED_STATUS = {401, 403, 451} +_RETRYABLE_STATUS_FLOOR = 500 + + +def utc_now_iso() -> str: + return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") + + +def classify_exception(exc: BaseException) -> tuple[str, bool]: + if isinstance(exc, (asyncio.TimeoutError, TimeoutError, httpx.TimeoutException)): + return "timeout", True + if isinstance(exc, httpx.HTTPStatusError): + code = exc.response.status_code + if code == 429: + return "rate_limited", True + if code in _BLOCKED_STATUS: + return "blocked", False + return "http_error", code >= _RETRYABLE_STATUS_FLOOR + if isinstance(exc, PermissionError): + return "no_privilege", False + if isinstance(exc, NotImplementedError): + return "not_applicable", False + if isinstance(exc, (httpx.TransportError, ConnectionError, OSError)): + return "unavailable", True + if isinstance(exc, (ValueError, KeyError, TypeError, IndexError, AttributeError)): + return "parse_error", False + return "unavailable", False + + +async def run_module( + name: str, + coro: Awaitable[Any], + *, + timeout: float, + source: str | None = None, +) -> ModuleResult: + started_at = utc_now_iso() + began = time.perf_counter() + try: + value = await asyncio.wait_for(coro, timeout=timeout) + except asyncio.CancelledError: + raise + except BaseException as exc: # noqa: BLE001 - failure is a value here, not control flow + kind, retryable = classify_exception(exc) + message = str(exc) or exc.__class__.__name__ + result = ModuleResult( + name=name, + status="failed", + errors=[ProbeError(source=source or name, kind=kind, message=message, retryable=retryable)], + ) + else: + if isinstance(value, ModuleResult): + result = value + else: + result = ModuleResult(name=name, status="ok", data=value) + result.name = name + result.started_at = started_at + result.duration_ms = int((time.perf_counter() - began) * 1000) + return result + + +async def gather_modules(*results: Awaitable[ModuleResult]) -> list[ModuleResult]: + gathered = await asyncio.gather(*results, return_exceptions=True) + out: list[ModuleResult] = [] + for item in gathered: + if isinstance(item, ModuleResult): + out.append(item) + continue + if isinstance(item, asyncio.CancelledError): + raise item + kind, retryable = classify_exception(item) + out.append( + ModuleResult( + name="unknown", + status="failed", + errors=[ProbeError(source="orchestration", kind=kind, message=str(item), retryable=retryable)], + ) + ) + return out diff --git a/tests/test_orchestration.py b/tests/test_orchestration.py new file mode 100644 index 0000000..7ff10bb --- /dev/null +++ b/tests/test_orchestration.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +import asyncio + +import httpx +import pytest + +from netcheck.models import ModuleResult, ProbeError +from netcheck.orchestration import classify_exception, gather_modules, run_module + + +def test_classify_timeout_errors(): + assert classify_exception(asyncio.TimeoutError()) == ("timeout", True) + assert classify_exception(httpx.ConnectTimeout("slow")) == ("timeout", True) + assert classify_exception(TimeoutError()) == ("timeout", True) + + +def test_classify_http_status_errors_by_code(): + def err(code: int) -> httpx.HTTPStatusError: + request = httpx.Request("GET", "https://example.test/") + response = httpx.Response(code, request=request) + return httpx.HTTPStatusError("boom", request=request, response=response) + + assert classify_exception(err(429)) == ("rate_limited", True) + assert classify_exception(err(403)) == ("blocked", False) + assert classify_exception(err(451)) == ("blocked", False) + assert classify_exception(err(500)) == ("http_error", True) + assert classify_exception(err(404)) == ("http_error", False) + + +def test_classify_transport_and_parse_and_privilege_errors(): + assert classify_exception(httpx.ConnectError("refused")) == ("unavailable", True) + assert classify_exception(OSError("network unreachable")) == ("unavailable", True) + assert classify_exception(ValueError("bad json")) == ("parse_error", False) + assert classify_exception(KeyError("asn")) == ("parse_error", False) + assert classify_exception(PermissionError("raw socket")) == ("no_privilege", False) + assert classify_exception(NotImplementedError("no win api here")) == ("not_applicable", False) + + +def test_classify_unknown_exception_falls_back_to_unavailable(): + class Weird(Exception): + pass + + assert classify_exception(Weird("?")) == ("unavailable", False) + + +async def test_run_module_returns_ok_envelope_on_success(): + async def work(): + await asyncio.sleep(0) + return {"value": 42} + + result = await run_module("bgp", work(), timeout=1.0) + assert isinstance(result, ModuleResult) + assert result.name == "bgp" + assert result.status == "ok" + assert result.data == {"value": 42} + assert result.errors == [] + assert result.started_at.endswith("Z") + assert result.duration_ms >= 0 + + +async def test_run_module_passes_through_a_module_result_unchanged_but_timed(): + async def work(): + return ModuleResult(name="reputation", status="partial", data={"x": 1}, warnings=["no key"]) + + result = await run_module("reputation", work(), timeout=1.0) + assert result.status == "partial" + assert result.warnings == ["no key"] + assert result.started_at.endswith("Z") + + +async def test_run_module_converts_an_exception_into_a_failed_envelope(): + async def work(): + raise httpx.ConnectError("refused") + + result = await run_module("ip_geo", work(), timeout=1.0, source="ip-api") + assert result.status == "failed" + assert result.data is None + assert result.errors == [ + ProbeError(source="ip-api", kind="unavailable", message="refused", retryable=True) + ] + + +async def test_run_module_enforces_its_own_timeout(): + async def work(): + await asyncio.sleep(5) + + result = await run_module("speed", work(), timeout=0.05) + assert result.status == "failed" + assert result.errors[0].kind == "timeout" + assert result.duration_ms < 2000 + + +async def test_run_module_lets_cancellation_propagate(): + async def work(): + raise asyncio.CancelledError() + + with pytest.raises(asyncio.CancelledError): + await run_module("latency", work(), timeout=1.0) + + +async def test_run_module_catches_base_exceptions_other_than_cancellation(): + # A literal KeyboardInterrupt can't be used here: CPython's asyncio.Task + # special-cases KeyboardInterrupt/SystemExit and re-raises them straight + # into the event loop's callback processing (in addition to setting them + # as the task's exception), which crashes asyncio.run() regardless of any + # `except BaseException` in the awaiting coroutine. A plain BaseException + # subclass exercises the same "not Exception, not CancelledError" path + # without hitting that special case. + class NonCancelBaseException(BaseException): + pass + + async def work(): + raise NonCancelBaseException() + + result = await run_module("dns_leak", work(), timeout=1.0) + assert result.status == "failed" + assert result.errors[0].kind == "unavailable" + + +async def test_gather_modules_never_raises_and_preserves_order(): + async def ok(): + return ModuleResult(name="a", status="ok", data=1) + + async def boom(): + raise RuntimeError("nope") + + results = await gather_modules( + run_module("a", ok(), timeout=1.0), + run_module("b", boom(), timeout=1.0), + ) + assert [r.name for r in results] == ["a", "b"] + assert [r.status for r in results] == ["ok", "failed"]