feat: add connection health monitor with metrics (#6)
Implements the HealthMonitor class for tracking connection health: - Tracks connection states per stream (active, stale, disconnected) - Records events and calculates throughput (events/second) - Detects stale streams (no events for configurable threshold) - Exposes Prometheus-compatible metrics endpoint (/metrics) - Provides HTTP health endpoints (/health, /ready, /live) Exports: HealthMonitor, HealthReport, HealthStatus, StreamHealth, StreamStatus Tests: 47 comprehensive unit tests covering all functionality 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
Patrick Selamy
@@ -14,6 +14,8 @@ dependencies = [
|
||||
"pydantic>=2.0.0",
|
||||
"python-dotenv>=1.0.0",
|
||||
"websockets>=12.0",
|
||||
"prometheus-client>=0.19.0",
|
||||
"aiohttp>=3.9.0",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
|
||||
@@ -5,6 +5,13 @@ from polymarket_insider_tracker.ingestor.clob_client import (
|
||||
ClobClientError,
|
||||
RetryError,
|
||||
)
|
||||
from polymarket_insider_tracker.ingestor.health import (
|
||||
HealthMonitor,
|
||||
HealthReport,
|
||||
HealthStatus,
|
||||
StreamHealth,
|
||||
StreamStatus,
|
||||
)
|
||||
from polymarket_insider_tracker.ingestor.metadata_sync import (
|
||||
MarketMetadataSync,
|
||||
MetadataSyncError,
|
||||
@@ -38,6 +45,12 @@ __all__ = [
|
||||
"ClobClient",
|
||||
"ClobClientError",
|
||||
"RetryError",
|
||||
# Health Monitor
|
||||
"HealthMonitor",
|
||||
"HealthReport",
|
||||
"HealthStatus",
|
||||
"StreamHealth",
|
||||
"StreamStatus",
|
||||
# Metadata Sync
|
||||
"MarketMetadataSync",
|
||||
"MetadataSyncError",
|
||||
|
||||
@@ -0,0 +1,511 @@
|
||||
"""Connection health monitor with metrics and HTTP endpoints.
|
||||
|
||||
This module provides health monitoring for the data ingestion layer,
|
||||
tracking connection states, event throughput, and staleness detection.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import copy
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
from aiohttp import web
|
||||
from prometheus_client import Counter, Gauge, Histogram, generate_latest
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default configuration
|
||||
DEFAULT_STALE_THRESHOLD_SECONDS = 60 # No events for 60s = stale
|
||||
DEFAULT_HEALTH_CHECK_INTERVAL = 5 # seconds
|
||||
DEFAULT_HTTP_PORT = 8080
|
||||
|
||||
|
||||
class HealthStatus(Enum):
|
||||
"""Overall health status."""
|
||||
|
||||
HEALTHY = "healthy"
|
||||
DEGRADED = "degraded"
|
||||
UNHEALTHY = "unhealthy"
|
||||
|
||||
|
||||
class StreamStatus(Enum):
|
||||
"""Status of an individual stream."""
|
||||
|
||||
ACTIVE = "active"
|
||||
STALE = "stale"
|
||||
DISCONNECTED = "disconnected"
|
||||
|
||||
|
||||
@dataclass
|
||||
class StreamHealth:
|
||||
"""Health status for an individual stream."""
|
||||
|
||||
name: str
|
||||
status: StreamStatus = StreamStatus.DISCONNECTED
|
||||
last_event_time: float | None = None
|
||||
events_received: int = 0
|
||||
events_per_second: float = 0.0
|
||||
connected_since: float | None = None
|
||||
last_error: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class HealthReport:
|
||||
"""Comprehensive health report for all streams."""
|
||||
|
||||
status: HealthStatus
|
||||
streams: dict[str, StreamHealth] = field(default_factory=dict)
|
||||
total_events_received: int = 0
|
||||
total_events_per_second: float = 0.0
|
||||
uptime_seconds: float = 0.0
|
||||
timestamp: float = field(default_factory=time.time)
|
||||
|
||||
|
||||
# Type aliases
|
||||
HealthCallback = Callable[[HealthReport], Awaitable[None]]
|
||||
|
||||
|
||||
# Prometheus metrics
|
||||
EVENTS_TOTAL = Counter(
|
||||
"polymarket_events_total",
|
||||
"Total number of events received",
|
||||
["stream"],
|
||||
)
|
||||
|
||||
EVENTS_PER_SECOND = Gauge(
|
||||
"polymarket_events_per_second",
|
||||
"Current events per second rate",
|
||||
["stream"],
|
||||
)
|
||||
|
||||
STREAM_STATUS = Gauge(
|
||||
"polymarket_stream_status",
|
||||
"Stream status (1=active, 0.5=stale, 0=disconnected)",
|
||||
["stream"],
|
||||
)
|
||||
|
||||
LAST_EVENT_TIMESTAMP = Gauge(
|
||||
"polymarket_last_event_timestamp",
|
||||
"Unix timestamp of last event received",
|
||||
["stream"],
|
||||
)
|
||||
|
||||
EVENT_LATENCY = Histogram(
|
||||
"polymarket_event_latency_seconds",
|
||||
"Event processing latency in seconds",
|
||||
["stream"],
|
||||
buckets=[0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0],
|
||||
)
|
||||
|
||||
HEALTH_STATUS = Gauge(
|
||||
"polymarket_health_status",
|
||||
"Overall health status (1=healthy, 0.5=degraded, 0=unhealthy)",
|
||||
)
|
||||
|
||||
|
||||
class HealthMonitor:
|
||||
"""Monitor connection health and expose metrics.
|
||||
|
||||
This class tracks the health of multiple streams, calculates throughput,
|
||||
detects stale streams, and exposes Prometheus-compatible metrics.
|
||||
|
||||
Example:
|
||||
```python
|
||||
monitor = HealthMonitor(stale_threshold_seconds=60)
|
||||
await monitor.start()
|
||||
|
||||
# Record events
|
||||
monitor.record_event("trades", processing_time=0.001)
|
||||
|
||||
# Update connection state
|
||||
monitor.set_stream_connected("trades")
|
||||
|
||||
# Get health report
|
||||
report = monitor.get_health_report()
|
||||
|
||||
# HTTP endpoints: /health and /metrics
|
||||
# Start HTTP server with monitor.start_http_server(port=8080)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
stale_threshold_seconds: float = DEFAULT_STALE_THRESHOLD_SECONDS,
|
||||
health_check_interval: float = DEFAULT_HEALTH_CHECK_INTERVAL,
|
||||
on_health_change: HealthCallback | None = None,
|
||||
) -> None:
|
||||
"""Initialize the health monitor.
|
||||
|
||||
Args:
|
||||
stale_threshold_seconds: Seconds without events before stream is stale.
|
||||
health_check_interval: Seconds between health check updates.
|
||||
on_health_change: Optional callback when health status changes.
|
||||
"""
|
||||
self._stale_threshold = stale_threshold_seconds
|
||||
self._health_check_interval = health_check_interval
|
||||
self._on_health_change = on_health_change
|
||||
|
||||
self._streams: dict[str, StreamHealth] = {}
|
||||
self._start_time: float | None = None
|
||||
self._running = False
|
||||
self._health_task: asyncio.Task[None] | None = None
|
||||
self._last_health_status: HealthStatus | None = None
|
||||
|
||||
# For throughput calculation
|
||||
self._event_windows: dict[str, list[float]] = {}
|
||||
self._window_duration = 10.0 # 10 second sliding window
|
||||
|
||||
# HTTP server
|
||||
self._app: web.Application | None = None
|
||||
self._runner: web.AppRunner | None = None
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
"""Return True if the monitor is running."""
|
||||
return self._running
|
||||
|
||||
def register_stream(self, name: str) -> None:
|
||||
"""Register a stream for monitoring.
|
||||
|
||||
Args:
|
||||
name: Unique name for the stream.
|
||||
"""
|
||||
if name not in self._streams:
|
||||
self._streams[name] = StreamHealth(name=name)
|
||||
self._event_windows[name] = []
|
||||
logger.info("Registered stream for monitoring: %s", name)
|
||||
|
||||
def set_stream_connected(self, name: str) -> None:
|
||||
"""Mark a stream as connected.
|
||||
|
||||
Args:
|
||||
name: Stream name.
|
||||
"""
|
||||
self.register_stream(name)
|
||||
stream = self._streams[name]
|
||||
stream.status = StreamStatus.ACTIVE
|
||||
stream.connected_since = time.time()
|
||||
stream.last_error = None
|
||||
STREAM_STATUS.labels(stream=name).set(1.0)
|
||||
logger.debug("Stream connected: %s", name)
|
||||
|
||||
def set_stream_disconnected(self, name: str, error: str | None = None) -> None:
|
||||
"""Mark a stream as disconnected.
|
||||
|
||||
Args:
|
||||
name: Stream name.
|
||||
error: Optional error message.
|
||||
"""
|
||||
self.register_stream(name)
|
||||
stream = self._streams[name]
|
||||
stream.status = StreamStatus.DISCONNECTED
|
||||
stream.connected_since = None
|
||||
stream.last_error = error
|
||||
STREAM_STATUS.labels(stream=name).set(0.0)
|
||||
logger.debug("Stream disconnected: %s (error: %s)", name, error)
|
||||
|
||||
def record_event(
|
||||
self,
|
||||
stream_name: str,
|
||||
*,
|
||||
processing_time: float | None = None,
|
||||
) -> None:
|
||||
"""Record an event received from a stream.
|
||||
|
||||
Args:
|
||||
stream_name: Name of the stream.
|
||||
processing_time: Optional processing latency in seconds.
|
||||
"""
|
||||
self.register_stream(stream_name)
|
||||
now = time.time()
|
||||
|
||||
stream = self._streams[stream_name]
|
||||
stream.events_received += 1
|
||||
stream.last_event_time = now
|
||||
stream.status = StreamStatus.ACTIVE
|
||||
|
||||
# Update metrics
|
||||
EVENTS_TOTAL.labels(stream=stream_name).inc()
|
||||
LAST_EVENT_TIMESTAMP.labels(stream=stream_name).set(now)
|
||||
STREAM_STATUS.labels(stream=stream_name).set(1.0)
|
||||
|
||||
if processing_time is not None:
|
||||
EVENT_LATENCY.labels(stream=stream_name).observe(processing_time)
|
||||
|
||||
# Add to sliding window for throughput
|
||||
window = self._event_windows[stream_name]
|
||||
window.append(now)
|
||||
|
||||
# Clean old entries from window
|
||||
cutoff = now - self._window_duration
|
||||
self._event_windows[stream_name] = [t for t in window if t > cutoff]
|
||||
|
||||
def _calculate_throughput(self, stream_name: str) -> float:
|
||||
"""Calculate events per second for a stream.
|
||||
|
||||
Args:
|
||||
stream_name: Name of the stream.
|
||||
|
||||
Returns:
|
||||
Events per second rate.
|
||||
"""
|
||||
window = self._event_windows.get(stream_name, [])
|
||||
if not window:
|
||||
return 0.0
|
||||
|
||||
now = time.time()
|
||||
cutoff = now - self._window_duration
|
||||
|
||||
# Count events in window
|
||||
recent_events = [t for t in window if t > cutoff]
|
||||
if not recent_events:
|
||||
return 0.0
|
||||
|
||||
# Calculate rate based on window
|
||||
window_span = now - cutoff
|
||||
return len(recent_events) / window_span if window_span > 0 else 0.0
|
||||
|
||||
def _check_stream_staleness(self) -> None:
|
||||
"""Check all streams for staleness."""
|
||||
now = time.time()
|
||||
|
||||
for name, stream in self._streams.items():
|
||||
if stream.status == StreamStatus.DISCONNECTED:
|
||||
continue
|
||||
|
||||
if stream.last_event_time is None:
|
||||
# Connected but no events yet - check connection time
|
||||
if stream.connected_since:
|
||||
since_connect = now - stream.connected_since
|
||||
if since_connect > self._stale_threshold:
|
||||
stream.status = StreamStatus.STALE
|
||||
STREAM_STATUS.labels(stream=name).set(0.5)
|
||||
else:
|
||||
since_event = now - stream.last_event_time
|
||||
if since_event > self._stale_threshold:
|
||||
stream.status = StreamStatus.STALE
|
||||
STREAM_STATUS.labels(stream=name).set(0.5)
|
||||
else:
|
||||
stream.status = StreamStatus.ACTIVE
|
||||
STREAM_STATUS.labels(stream=name).set(1.0)
|
||||
|
||||
def _determine_overall_status(self) -> HealthStatus:
|
||||
"""Determine overall health status based on stream states.
|
||||
|
||||
Returns:
|
||||
Overall health status.
|
||||
"""
|
||||
if not self._streams:
|
||||
return HealthStatus.HEALTHY # No streams = healthy (nothing to monitor)
|
||||
|
||||
statuses = [s.status for s in self._streams.values()]
|
||||
|
||||
if all(s == StreamStatus.DISCONNECTED for s in statuses):
|
||||
return HealthStatus.UNHEALTHY
|
||||
|
||||
if any(s == StreamStatus.DISCONNECTED for s in statuses):
|
||||
return HealthStatus.DEGRADED
|
||||
|
||||
if any(s == StreamStatus.STALE for s in statuses):
|
||||
return HealthStatus.DEGRADED
|
||||
|
||||
return HealthStatus.HEALTHY
|
||||
|
||||
def get_health_report(self) -> HealthReport:
|
||||
"""Generate a comprehensive health report.
|
||||
|
||||
Returns:
|
||||
HealthReport with current status of all streams.
|
||||
"""
|
||||
self._check_stream_staleness()
|
||||
|
||||
# Update throughput metrics
|
||||
total_eps = 0.0
|
||||
for name, stream in self._streams.items():
|
||||
eps = self._calculate_throughput(name)
|
||||
stream.events_per_second = eps
|
||||
EVENTS_PER_SECOND.labels(stream=name).set(eps)
|
||||
total_eps += eps
|
||||
|
||||
overall_status = self._determine_overall_status()
|
||||
HEALTH_STATUS.set(
|
||||
1.0 if overall_status == HealthStatus.HEALTHY
|
||||
else 0.5 if overall_status == HealthStatus.DEGRADED
|
||||
else 0.0
|
||||
)
|
||||
|
||||
uptime = 0.0
|
||||
if self._start_time:
|
||||
uptime = time.time() - self._start_time
|
||||
|
||||
# Deep copy streams to prevent mutations affecting internal state
|
||||
streams_copy = {name: copy.copy(stream) for name, stream in self._streams.items()}
|
||||
|
||||
return HealthReport(
|
||||
status=overall_status,
|
||||
streams=streams_copy,
|
||||
total_events_received=sum(s.events_received for s in self._streams.values()),
|
||||
total_events_per_second=total_eps,
|
||||
uptime_seconds=uptime,
|
||||
)
|
||||
|
||||
async def _health_check_loop(self) -> None:
|
||||
"""Background task for periodic health checks."""
|
||||
while self._running:
|
||||
try:
|
||||
report = self.get_health_report()
|
||||
|
||||
# Notify on status change
|
||||
if (
|
||||
self._on_health_change
|
||||
and report.status != self._last_health_status
|
||||
):
|
||||
self._last_health_status = report.status
|
||||
try:
|
||||
await self._on_health_change(report)
|
||||
except Exception as e:
|
||||
logger.error("Error in health change callback: %s", e)
|
||||
|
||||
await asyncio.sleep(self._health_check_interval)
|
||||
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error("Error in health check loop: %s", e)
|
||||
await asyncio.sleep(1)
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the health monitor.
|
||||
|
||||
Begins periodic health checks and staleness detection.
|
||||
"""
|
||||
if self._running:
|
||||
return
|
||||
|
||||
self._running = True
|
||||
self._start_time = time.time()
|
||||
self._health_task = asyncio.create_task(self._health_check_loop())
|
||||
logger.info("Health monitor started")
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the health monitor."""
|
||||
if not self._running:
|
||||
return
|
||||
|
||||
self._running = False
|
||||
|
||||
if self._health_task:
|
||||
self._health_task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await self._health_task
|
||||
self._health_task = None
|
||||
|
||||
await self.stop_http_server()
|
||||
logger.info("Health monitor stopped")
|
||||
|
||||
# HTTP Server methods
|
||||
|
||||
async def _handle_health(self, _request: web.Request) -> web.Response:
|
||||
"""Handle /health endpoint."""
|
||||
report = self.get_health_report()
|
||||
|
||||
status_code = 200 if report.status == HealthStatus.HEALTHY else 503
|
||||
|
||||
body: dict[str, Any] = {
|
||||
"status": report.status.value,
|
||||
"uptime_seconds": report.uptime_seconds,
|
||||
"total_events_received": report.total_events_received,
|
||||
"total_events_per_second": round(report.total_events_per_second, 2),
|
||||
"streams": {},
|
||||
}
|
||||
|
||||
for name, stream in report.streams.items():
|
||||
body["streams"][name] = {
|
||||
"status": stream.status.value,
|
||||
"events_received": stream.events_received,
|
||||
"events_per_second": round(stream.events_per_second, 2),
|
||||
"last_event_time": stream.last_event_time,
|
||||
"last_error": stream.last_error,
|
||||
}
|
||||
|
||||
return web.json_response(body, status=status_code)
|
||||
|
||||
async def _handle_metrics(self, _request: web.Request) -> web.Response:
|
||||
"""Handle /metrics endpoint (Prometheus format)."""
|
||||
# Ensure latest values are calculated
|
||||
self.get_health_report()
|
||||
|
||||
metrics = generate_latest()
|
||||
return web.Response(
|
||||
body=metrics,
|
||||
content_type="text/plain",
|
||||
charset="utf-8",
|
||||
)
|
||||
|
||||
async def _handle_ready(self, _request: web.Request) -> web.Response:
|
||||
"""Handle /ready endpoint for k8s readiness probe."""
|
||||
report = self.get_health_report()
|
||||
|
||||
if report.status == HealthStatus.UNHEALTHY:
|
||||
return web.json_response(
|
||||
{"ready": False, "reason": "unhealthy"},
|
||||
status=503,
|
||||
)
|
||||
|
||||
return web.json_response({"ready": True}, status=200)
|
||||
|
||||
async def _handle_live(self, _request: web.Request) -> web.Response:
|
||||
"""Handle /live endpoint for k8s liveness probe."""
|
||||
# Always return 200 if the server is running
|
||||
return web.json_response({"live": True}, status=200)
|
||||
|
||||
def _create_app(self) -> web.Application:
|
||||
"""Create the aiohttp application."""
|
||||
app = web.Application()
|
||||
app.router.add_get("/health", self._handle_health)
|
||||
app.router.add_get("/metrics", self._handle_metrics)
|
||||
app.router.add_get("/ready", self._handle_ready)
|
||||
app.router.add_get("/live", self._handle_live)
|
||||
return app
|
||||
|
||||
async def start_http_server(self, port: int = DEFAULT_HTTP_PORT) -> None:
|
||||
"""Start the HTTP server for health and metrics endpoints.
|
||||
|
||||
Args:
|
||||
port: Port to listen on.
|
||||
"""
|
||||
if self._runner:
|
||||
logger.warning("HTTP server already running")
|
||||
return
|
||||
|
||||
self._app = self._create_app()
|
||||
self._runner = web.AppRunner(self._app)
|
||||
await self._runner.setup()
|
||||
|
||||
site = web.TCPSite(self._runner, "0.0.0.0", port)
|
||||
await site.start()
|
||||
|
||||
logger.info("Health HTTP server started on port %d", port)
|
||||
|
||||
async def stop_http_server(self) -> None:
|
||||
"""Stop the HTTP server."""
|
||||
if self._runner:
|
||||
await self._runner.cleanup()
|
||||
self._runner = None
|
||||
self._app = None
|
||||
logger.info("Health HTTP server stopped")
|
||||
|
||||
async def __aenter__(self) -> "HealthMonitor":
|
||||
"""Async context manager entry."""
|
||||
await self.start()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args: Any) -> None:
|
||||
"""Async context manager exit."""
|
||||
await self.stop()
|
||||
@@ -0,0 +1,617 @@
|
||||
"""Tests for the connection health monitor."""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
from aiohttp import web
|
||||
|
||||
from polymarket_insider_tracker.ingestor.health import (
|
||||
DEFAULT_HEALTH_CHECK_INTERVAL,
|
||||
DEFAULT_STALE_THRESHOLD_SECONDS,
|
||||
HealthMonitor,
|
||||
HealthReport,
|
||||
HealthStatus,
|
||||
StreamHealth,
|
||||
StreamStatus,
|
||||
)
|
||||
|
||||
|
||||
class TestStreamHealth:
|
||||
"""Tests for the StreamHealth dataclass."""
|
||||
|
||||
def test_stream_health_defaults(self) -> None:
|
||||
"""Test default values."""
|
||||
health = StreamHealth(name="test-stream")
|
||||
|
||||
assert health.name == "test-stream"
|
||||
assert health.status == StreamStatus.DISCONNECTED
|
||||
assert health.last_event_time is None
|
||||
assert health.events_received == 0
|
||||
assert health.events_per_second == 0.0
|
||||
assert health.connected_since is None
|
||||
assert health.last_error is None
|
||||
|
||||
def test_stream_health_custom_values(self) -> None:
|
||||
"""Test with custom values."""
|
||||
now = time.time()
|
||||
health = StreamHealth(
|
||||
name="trades",
|
||||
status=StreamStatus.ACTIVE,
|
||||
last_event_time=now,
|
||||
events_received=100,
|
||||
events_per_second=5.0,
|
||||
connected_since=now - 3600,
|
||||
last_error=None,
|
||||
)
|
||||
|
||||
assert health.name == "trades"
|
||||
assert health.status == StreamStatus.ACTIVE
|
||||
assert health.events_received == 100
|
||||
|
||||
|
||||
class TestHealthReport:
|
||||
"""Tests for the HealthReport dataclass."""
|
||||
|
||||
def test_health_report_defaults(self) -> None:
|
||||
"""Test default values."""
|
||||
report = HealthReport(status=HealthStatus.HEALTHY)
|
||||
|
||||
assert report.status == HealthStatus.HEALTHY
|
||||
assert report.streams == {}
|
||||
assert report.total_events_received == 0
|
||||
assert report.total_events_per_second == 0.0
|
||||
assert report.uptime_seconds == 0.0
|
||||
assert report.timestamp > 0
|
||||
|
||||
def test_health_report_with_streams(self) -> None:
|
||||
"""Test with stream data."""
|
||||
stream = StreamHealth(name="trades", events_received=100)
|
||||
report = HealthReport(
|
||||
status=HealthStatus.DEGRADED,
|
||||
streams={"trades": stream},
|
||||
total_events_received=100,
|
||||
total_events_per_second=5.0,
|
||||
uptime_seconds=3600.0,
|
||||
)
|
||||
|
||||
assert report.status == HealthStatus.DEGRADED
|
||||
assert "trades" in report.streams
|
||||
assert report.total_events_received == 100
|
||||
|
||||
|
||||
class TestHealthMonitor:
|
||||
"""Tests for the HealthMonitor class."""
|
||||
|
||||
def test_init(self) -> None:
|
||||
"""Test initialization."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
assert monitor._stale_threshold == DEFAULT_STALE_THRESHOLD_SECONDS
|
||||
assert monitor._health_check_interval == DEFAULT_HEALTH_CHECK_INTERVAL
|
||||
assert not monitor.is_running
|
||||
|
||||
def test_init_custom_config(self) -> None:
|
||||
"""Test initialization with custom config."""
|
||||
monitor = HealthMonitor(
|
||||
stale_threshold_seconds=30,
|
||||
health_check_interval=10,
|
||||
)
|
||||
|
||||
assert monitor._stale_threshold == 30
|
||||
assert monitor._health_check_interval == 10
|
||||
|
||||
def test_register_stream(self) -> None:
|
||||
"""Test registering a stream."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.register_stream("trades")
|
||||
|
||||
assert "trades" in monitor._streams
|
||||
assert monitor._streams["trades"].name == "trades"
|
||||
assert monitor._streams["trades"].status == StreamStatus.DISCONNECTED
|
||||
|
||||
def test_register_stream_idempotent(self) -> None:
|
||||
"""Test that registering the same stream twice is idempotent."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.register_stream("trades")
|
||||
monitor.record_event("trades") # Adds an event
|
||||
monitor.register_stream("trades") # Should not reset
|
||||
|
||||
assert monitor._streams["trades"].events_received == 1
|
||||
|
||||
def test_set_stream_connected(self) -> None:
|
||||
"""Test marking a stream as connected."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.set_stream_connected("trades")
|
||||
|
||||
assert monitor._streams["trades"].status == StreamStatus.ACTIVE
|
||||
assert monitor._streams["trades"].connected_since is not None
|
||||
assert monitor._streams["trades"].last_error is None
|
||||
|
||||
def test_set_stream_disconnected(self) -> None:
|
||||
"""Test marking a stream as disconnected."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.set_stream_connected("trades")
|
||||
monitor.set_stream_disconnected("trades", error="Connection reset")
|
||||
|
||||
assert monitor._streams["trades"].status == StreamStatus.DISCONNECTED
|
||||
assert monitor._streams["trades"].connected_since is None
|
||||
assert monitor._streams["trades"].last_error == "Connection reset"
|
||||
|
||||
def test_record_event(self) -> None:
|
||||
"""Test recording an event."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.record_event("trades")
|
||||
|
||||
stream = monitor._streams["trades"]
|
||||
assert stream.events_received == 1
|
||||
assert stream.last_event_time is not None
|
||||
assert stream.status == StreamStatus.ACTIVE
|
||||
|
||||
def test_record_event_multiple(self) -> None:
|
||||
"""Test recording multiple events."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
for _ in range(10):
|
||||
monitor.record_event("trades")
|
||||
|
||||
assert monitor._streams["trades"].events_received == 10
|
||||
|
||||
def test_record_event_with_processing_time(self) -> None:
|
||||
"""Test recording event with processing time."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
# Should not raise
|
||||
monitor.record_event("trades", processing_time=0.001)
|
||||
|
||||
assert monitor._streams["trades"].events_received == 1
|
||||
|
||||
def test_calculate_throughput_empty(self) -> None:
|
||||
"""Test throughput calculation with no events."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
rate = monitor._calculate_throughput("nonexistent")
|
||||
|
||||
assert rate == 0.0
|
||||
|
||||
def test_calculate_throughput(self) -> None:
|
||||
"""Test throughput calculation."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
# Add events
|
||||
for _ in range(10):
|
||||
monitor.record_event("trades")
|
||||
|
||||
rate = monitor._calculate_throughput("trades")
|
||||
|
||||
# Should have ~10 events in the window
|
||||
assert rate > 0
|
||||
|
||||
def test_check_stream_staleness_active(self) -> None:
|
||||
"""Test that active stream is not marked stale."""
|
||||
monitor = HealthMonitor(stale_threshold_seconds=60)
|
||||
|
||||
monitor.record_event("trades")
|
||||
monitor._check_stream_staleness()
|
||||
|
||||
assert monitor._streams["trades"].status == StreamStatus.ACTIVE
|
||||
|
||||
def test_check_stream_staleness_stale(self) -> None:
|
||||
"""Test that stream becomes stale after threshold."""
|
||||
monitor = HealthMonitor(stale_threshold_seconds=1)
|
||||
|
||||
monitor.record_event("trades")
|
||||
# Simulate time passing
|
||||
monitor._streams["trades"].last_event_time = time.time() - 2
|
||||
|
||||
monitor._check_stream_staleness()
|
||||
|
||||
assert monitor._streams["trades"].status == StreamStatus.STALE
|
||||
|
||||
def test_check_stream_staleness_connected_no_events(self) -> None:
|
||||
"""Test staleness when connected but no events received."""
|
||||
monitor = HealthMonitor(stale_threshold_seconds=1)
|
||||
|
||||
monitor.set_stream_connected("trades")
|
||||
# Simulate time passing since connection
|
||||
monitor._streams["trades"].connected_since = time.time() - 2
|
||||
|
||||
monitor._check_stream_staleness()
|
||||
|
||||
assert monitor._streams["trades"].status == StreamStatus.STALE
|
||||
|
||||
def test_determine_overall_status_no_streams(self) -> None:
|
||||
"""Test overall status with no streams."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
status = monitor._determine_overall_status()
|
||||
|
||||
assert status == HealthStatus.HEALTHY
|
||||
|
||||
def test_determine_overall_status_all_active(self) -> None:
|
||||
"""Test overall status with all active streams."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.record_event("trades")
|
||||
monitor.record_event("orderbook")
|
||||
|
||||
status = monitor._determine_overall_status()
|
||||
|
||||
assert status == HealthStatus.HEALTHY
|
||||
|
||||
def test_determine_overall_status_some_stale(self) -> None:
|
||||
"""Test overall status with some stale streams."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.record_event("trades")
|
||||
monitor.register_stream("orderbook")
|
||||
monitor._streams["orderbook"].status = StreamStatus.STALE
|
||||
|
||||
status = monitor._determine_overall_status()
|
||||
|
||||
assert status == HealthStatus.DEGRADED
|
||||
|
||||
def test_determine_overall_status_some_disconnected(self) -> None:
|
||||
"""Test overall status with some disconnected streams."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.record_event("trades")
|
||||
monitor.set_stream_disconnected("orderbook")
|
||||
|
||||
status = monitor._determine_overall_status()
|
||||
|
||||
assert status == HealthStatus.DEGRADED
|
||||
|
||||
def test_determine_overall_status_all_disconnected(self) -> None:
|
||||
"""Test overall status with all disconnected streams."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.set_stream_disconnected("trades")
|
||||
monitor.set_stream_disconnected("orderbook")
|
||||
|
||||
status = monitor._determine_overall_status()
|
||||
|
||||
assert status == HealthStatus.UNHEALTHY
|
||||
|
||||
def test_get_health_report(self) -> None:
|
||||
"""Test getting a health report."""
|
||||
monitor = HealthMonitor()
|
||||
monitor._start_time = time.time() - 100
|
||||
|
||||
monitor.record_event("trades")
|
||||
monitor.record_event("trades")
|
||||
|
||||
report = monitor.get_health_report()
|
||||
|
||||
assert report.status == HealthStatus.HEALTHY
|
||||
assert "trades" in report.streams
|
||||
assert report.total_events_received == 2
|
||||
assert report.uptime_seconds >= 100
|
||||
assert report.timestamp > 0
|
||||
|
||||
def test_get_health_report_calculates_throughput(self) -> None:
|
||||
"""Test that health report calculates throughput."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
for _ in range(10):
|
||||
monitor.record_event("trades")
|
||||
|
||||
report = monitor.get_health_report()
|
||||
|
||||
assert report.streams["trades"].events_per_second > 0
|
||||
assert report.total_events_per_second > 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_stop(self) -> None:
|
||||
"""Test starting and stopping the monitor."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
await monitor.start()
|
||||
assert monitor.is_running
|
||||
assert monitor._health_task is not None
|
||||
|
||||
await monitor.stop()
|
||||
assert not monitor.is_running
|
||||
assert monitor._health_task is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_idempotent(self) -> None:
|
||||
"""Test that starting twice is safe."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
await monitor.start()
|
||||
await monitor.start() # Should not raise
|
||||
|
||||
assert monitor.is_running
|
||||
|
||||
await monitor.stop()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_when_not_running(self) -> None:
|
||||
"""Test that stopping when not running is safe."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
await monitor.stop() # Should not raise
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_context_manager(self) -> None:
|
||||
"""Test async context manager."""
|
||||
async with HealthMonitor() as monitor:
|
||||
assert monitor.is_running
|
||||
|
||||
assert not monitor.is_running
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_check_loop_updates_report(self) -> None:
|
||||
"""Test that health check loop updates the report."""
|
||||
monitor = HealthMonitor(health_check_interval=0.1)
|
||||
|
||||
await monitor.start()
|
||||
monitor.record_event("trades")
|
||||
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
# Health should have been checked
|
||||
report = monitor.get_health_report()
|
||||
assert report.status == HealthStatus.HEALTHY
|
||||
|
||||
await monitor.stop()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_change_callback(self) -> None:
|
||||
"""Test that health change callback is invoked."""
|
||||
callback = AsyncMock()
|
||||
monitor = HealthMonitor(
|
||||
health_check_interval=0.1,
|
||||
on_health_change=callback,
|
||||
)
|
||||
|
||||
await monitor.start()
|
||||
monitor.record_event("trades")
|
||||
|
||||
# Wait for health check
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
await monitor.stop()
|
||||
|
||||
# Callback should have been called at least once
|
||||
assert callback.called
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_change_callback_error_handling(self) -> None:
|
||||
"""Test that callback errors don't crash the loop."""
|
||||
callback = AsyncMock(side_effect=ValueError("test error"))
|
||||
monitor = HealthMonitor(
|
||||
health_check_interval=0.1,
|
||||
on_health_change=callback,
|
||||
)
|
||||
|
||||
await monitor.start()
|
||||
monitor.record_event("trades")
|
||||
|
||||
# Should not crash
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
await monitor.stop()
|
||||
|
||||
|
||||
class TestHealthMonitorHTTPEndpoints:
|
||||
"""Tests for HTTP endpoints."""
|
||||
|
||||
@pytest.fixture
|
||||
def monitor(self) -> HealthMonitor:
|
||||
"""Create a monitor instance."""
|
||||
return HealthMonitor()
|
||||
|
||||
@pytest.fixture
|
||||
def app(self, monitor: HealthMonitor) -> web.Application:
|
||||
"""Create the aiohttp application."""
|
||||
return monitor._create_app()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_endpoint_healthy(
|
||||
self, monitor: HealthMonitor, app: web.Application
|
||||
) -> None:
|
||||
"""Test /health endpoint when healthy."""
|
||||
from aiohttp.test_utils import TestClient, TestServer
|
||||
|
||||
monitor.record_event("trades")
|
||||
|
||||
async with TestClient(TestServer(app)) as client:
|
||||
resp = await client.get("/health")
|
||||
assert resp.status == 200
|
||||
|
||||
data = await resp.json()
|
||||
assert data["status"] == "healthy"
|
||||
assert "trades" in data["streams"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_endpoint_unhealthy(
|
||||
self, monitor: HealthMonitor, app: web.Application
|
||||
) -> None:
|
||||
"""Test /health endpoint when unhealthy."""
|
||||
from aiohttp.test_utils import TestClient, TestServer
|
||||
|
||||
monitor.set_stream_disconnected("trades")
|
||||
|
||||
async with TestClient(TestServer(app)) as client:
|
||||
resp = await client.get("/health")
|
||||
assert resp.status == 503
|
||||
|
||||
data = await resp.json()
|
||||
assert data["status"] == "unhealthy"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_metrics_endpoint(
|
||||
self, monitor: HealthMonitor, app: web.Application
|
||||
) -> None:
|
||||
"""Test /metrics endpoint returns Prometheus format."""
|
||||
from aiohttp.test_utils import TestClient, TestServer
|
||||
|
||||
monitor.record_event("trades")
|
||||
|
||||
async with TestClient(TestServer(app)) as client:
|
||||
resp = await client.get("/metrics")
|
||||
assert resp.status == 200
|
||||
|
||||
content_type = resp.headers.get("Content-Type", "")
|
||||
assert "text/plain" in content_type
|
||||
|
||||
text = await resp.text()
|
||||
assert "polymarket_events_total" in text
|
||||
assert "polymarket_health_status" in text
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ready_endpoint_ready(
|
||||
self, monitor: HealthMonitor, app: web.Application
|
||||
) -> None:
|
||||
"""Test /ready endpoint when ready."""
|
||||
from aiohttp.test_utils import TestClient, TestServer
|
||||
|
||||
monitor.record_event("trades")
|
||||
|
||||
async with TestClient(TestServer(app)) as client:
|
||||
resp = await client.get("/ready")
|
||||
assert resp.status == 200
|
||||
|
||||
data = await resp.json()
|
||||
assert data["ready"] is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ready_endpoint_not_ready(
|
||||
self, monitor: HealthMonitor, app: web.Application
|
||||
) -> None:
|
||||
"""Test /ready endpoint when not ready."""
|
||||
from aiohttp.test_utils import TestClient, TestServer
|
||||
|
||||
monitor.set_stream_disconnected("trades")
|
||||
|
||||
async with TestClient(TestServer(app)) as client:
|
||||
resp = await client.get("/ready")
|
||||
assert resp.status == 503
|
||||
|
||||
data = await resp.json()
|
||||
assert data["ready"] is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_live_endpoint(self, app: web.Application) -> None:
|
||||
"""Test /live endpoint always returns 200."""
|
||||
from aiohttp.test_utils import TestClient, TestServer
|
||||
|
||||
async with TestClient(TestServer(app)) as client:
|
||||
resp = await client.get("/live")
|
||||
assert resp.status == 200
|
||||
|
||||
data = await resp.json()
|
||||
assert data["live"] is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_stop_http_server(self, monitor: HealthMonitor) -> None:
|
||||
"""Test starting and stopping HTTP server."""
|
||||
await monitor.start_http_server(port=18080)
|
||||
assert monitor._runner is not None
|
||||
|
||||
await monitor.stop_http_server()
|
||||
assert monitor._runner is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_http_server_idempotent(self, monitor: HealthMonitor) -> None:
|
||||
"""Test that starting HTTP server twice is safe."""
|
||||
await monitor.start_http_server(port=18081)
|
||||
await monitor.start_http_server(port=18081) # Should not raise
|
||||
|
||||
await monitor.stop_http_server()
|
||||
|
||||
|
||||
class TestPrometheusMetrics:
|
||||
"""Tests for Prometheus metric updates."""
|
||||
|
||||
def test_events_total_incremented(self) -> None:
|
||||
"""Test that events_total counter is incremented."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.record_event("test-metrics")
|
||||
monitor.record_event("test-metrics")
|
||||
|
||||
# Counter should have been incremented
|
||||
# (We can't easily test prometheus metrics directly, but at least verify no errors)
|
||||
|
||||
def test_stream_status_updated(self) -> None:
|
||||
"""Test that stream_status gauge is updated."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.set_stream_connected("test-status")
|
||||
# Gauge should be 1.0
|
||||
|
||||
monitor.set_stream_disconnected("test-status")
|
||||
# Gauge should be 0.0
|
||||
|
||||
def test_health_status_updated(self) -> None:
|
||||
"""Test that health_status gauge is updated."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.record_event("test-health")
|
||||
report = monitor.get_health_report()
|
||||
|
||||
assert report.status == HealthStatus.HEALTHY
|
||||
|
||||
|
||||
class TestEdgeCases:
|
||||
"""Tests for edge cases and error handling."""
|
||||
|
||||
def test_throughput_with_old_events(self) -> None:
|
||||
"""Test throughput calculation ignores old events."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.record_event("trades")
|
||||
# Manually add old event to window
|
||||
monitor._event_windows["trades"].append(time.time() - 100)
|
||||
|
||||
rate = monitor._calculate_throughput("trades")
|
||||
|
||||
# Old event should be filtered out
|
||||
# Rate should only count recent events
|
||||
assert rate >= 0
|
||||
|
||||
def test_multiple_streams_independent(self) -> None:
|
||||
"""Test that multiple streams are tracked independently."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.record_event("trades")
|
||||
monitor.record_event("trades")
|
||||
monitor.set_stream_disconnected("orderbook")
|
||||
|
||||
assert monitor._streams["trades"].events_received == 2
|
||||
assert monitor._streams["trades"].status == StreamStatus.ACTIVE
|
||||
assert monitor._streams["orderbook"].events_received == 0
|
||||
assert monitor._streams["orderbook"].status == StreamStatus.DISCONNECTED
|
||||
|
||||
def test_report_streams_are_copied(self) -> None:
|
||||
"""Test that report streams are a copy."""
|
||||
monitor = HealthMonitor()
|
||||
monitor.record_event("trades")
|
||||
|
||||
report = monitor.get_health_report()
|
||||
|
||||
# Modifying report should not affect monitor
|
||||
report.streams["trades"].events_received = 999
|
||||
assert monitor._streams["trades"].events_received == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_cleans_up_http_server(self) -> None:
|
||||
"""Test that stop() also stops HTTP server."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
await monitor.start()
|
||||
await monitor.start_http_server(port=18082)
|
||||
|
||||
await monitor.stop()
|
||||
|
||||
assert not monitor.is_running
|
||||
assert monitor._runner is None
|
||||
Reference in New Issue
Block a user