stonks-oracle/tests/intelligence_pipeline_v3/test_observability.py

"""Tests for observability module — Task 43: traces, metrics, alerts."""

from __future__ import annotations

from uuid import uuid4

from services.intelligence_pipeline_v3.observability.metrics import (
    AlertSeverity,
    MetricAlert,
    MetricsCollector,
    StageMetrics,
)
from services.intelligence_pipeline_v3.observability.tracing import (
    PipelineTrace,
    SpanStatus,
    TraceCollector,
)


class TestPipelineTracing:
    """Task 43.1: Trace every stage under one document trace ID."""

    def test_trace_creation(self):
        trace = PipelineTrace.create("doc-001", uuid4())
        assert trace.document_id == "doc-001"
        assert trace.trace_id is not None
        assert not trace.is_complete

    def test_start_and_finish_span(self):
        trace = PipelineTrace.create("doc-001", uuid4())
        span = trace.start_span("extraction")
        assert span.stage_name == "extraction"
        assert span.status == SpanStatus.RUNNING
        span.finish(SpanStatus.SUCCEEDED)
        assert span.status == SpanStatus.SUCCEEDED
        assert span.duration_ms >= 0

    def test_multiple_spans(self):
        trace = PipelineTrace.create("doc-001", uuid4())
        trace.start_span("segmentation").finish()
        trace.start_span("extraction").finish()
        trace.start_span("routing").finish()
        assert len(trace.spans) == 3

    def test_failed_spans_tracked(self):
        trace = PipelineTrace.create("doc-001", uuid4())
        trace.start_span("extraction").finish(SpanStatus.FAILED, "timeout")
        trace.start_span("routing").finish(SpanStatus.SUCCEEDED)
        assert len(trace.failed_spans) == 1

    def test_trace_finish(self):
        trace = PipelineTrace.create("doc-001", uuid4())
        trace.finish()
        assert trace.is_complete
        assert trace.total_duration_ms >= 0

    def test_to_dict_serialization(self):
        trace = PipelineTrace.create("doc-001", uuid4())
        trace.start_span("extraction").finish()
        trace.finish()
        d = trace.to_dict()
        assert d["document_id"] == "doc-001"
        assert d["span_count"] == 1
        assert "spans" in d


class TestTraceCollector:
    """Task 43.1: Trace collection and retrieval."""

    def test_start_and_get_trace(self):
        collector = TraceCollector()
        trace = collector.start_trace("doc-001", uuid4())
        retrieved = collector.get_trace(trace.trace_id)
        assert retrieved is trace

    def test_get_by_document(self):
        collector = TraceCollector()
        run1 = uuid4()
        run2 = uuid4()
        collector.start_trace("doc-001", run1)
        collector.start_trace("doc-001", run2)
        collector.start_trace("doc-002", uuid4())
        results = collector.get_by_document("doc-001")
        assert len(results) == 2

    def test_eviction_at_max(self):
        collector = TraceCollector(max_stored=3)
        for i in range(5):
            collector.start_trace(f"doc-{i}", uuid4())
        assert collector.trace_count == 3


class TestStageMetrics:
    """Task 43.2: Stage latency, errors, batch size, queue depth, routing."""

    def test_record_invocation(self):
        metrics = StageMetrics(stage_name="extraction")
        metrics.record_invocation(latency_ms=150.0, tokens_in=500, tokens_out=200)
        assert metrics.total_invocations == 1
        assert metrics.avg_latency_ms == 150.0
        assert metrics.error_rate == 0.0

    def test_error_rate(self):
        metrics = StageMetrics(stage_name="adjudication")
        metrics.record_invocation(latency_ms=100, error=True)
        metrics.record_invocation(latency_ms=100, error=False)
        assert metrics.error_rate == 0.5

    def test_gpu_metrics(self):
        metrics = StageMetrics(stage_name="adjudication")
        metrics.record_invocation(
            latency_ms=500, gpu_seconds=0.5, gpu_memory_mb=4096
        )
        assert metrics.gpu_seconds_per_doc == 0.5
        assert metrics.gpu_memory_peak_mb == 4096

    def test_batch_size_tracking(self):
        metrics = StageMetrics(stage_name="specialist")
        metrics.record_invocation(latency_ms=50, batch_size=8)
        metrics.record_invocation(latency_ms=50, batch_size=4)
        assert metrics.avg_batch_size == 6.0


class TestMetricsCollector:
    """Task 43.2-43.5: Metrics collection and alerts."""

    def test_record_stage(self):
        collector = MetricsCollector()
        collector.record_stage("extraction", latency_ms=100)
        stage = collector.get_stage("extraction")
        assert stage.total_invocations == 1

    def test_increment_counter(self):
        collector = MetricsCollector()
        collector.increment_counter("schema_failures", 3)
        assert collector.get_counter("schema_failures") == 3

    def test_alert_evaluation(self):
        alert = MetricAlert(
            name="test_alert",
            metric_name="error_rate",
            condition="> 0.05",
            severity=AlertSeverity.CRITICAL,
            description="Error rate high",
            threshold=0.05,
        )
        assert alert.evaluate(0.10)  # Should fire
        assert not alert.evaluate(0.03)  # Should not fire

    def test_check_alerts(self):
        collector = MetricsCollector()
        collector.increment_counter("schema_failures", 0.10)
        fired = collector.check_alerts()
        # schema_failure_rate_high should fire (0.10 > 0.05)
        assert any(a.name == "schema_failure_rate_high" for a, _ in fired)

    def test_summary(self):
        collector = MetricsCollector()
        collector.record_stage("extraction", latency_ms=100)
        summary = collector.summary()
        assert "stages" in summary
        assert "extraction" in summary["stages"]