tests/observability/test_operation.py

"""Behavioural tests for the single operation boundary.

These assert on the actual emitted artefacts -- local-store rows, span
attributes, recorded histogram points, correlated log records -- not merely
that no exception was raised.
"""

from __future__ import annotations

import logging

import pytest

from app.observability import bootstrap
from app.observability.config import ObservabilityConfig
from app.observability.contracts import ExperimentIdentity, RetrievalOutcome
from app.observability.local_store import LocalObservationStore
from app.observability.logging import emit_correlated_log
from app.observability.operation import observe_operation, record_event


@pytest.fixture(autouse=True)
def _shutdown_after(tmp_path):
    """Ensure global observability state is torn down between tests."""
    yield
    bootstrap.shutdown_observability(timeout_seconds=2.0)


def _local_config(tmp_path) -> ObservabilityConfig:
    return ObservabilityConfig(
        enabled=True,
        mode="local",
        otlp_endpoint="http://127.0.0.1:4318",
        signoz_ui_url="http://127.0.0.1:8080",
        artifact_root=tmp_path,
    )


# --- Brief Step 2 example ---------------------------------------------------


def test_local_mode_writes_the_same_completed_observation_used_by_span(tmp_path):
    store = LocalObservationStore(tmp_path)
    with observe_operation("retrieval.vector_search", store=store) as op:
        op.set("raw_candidate_count", 12)
    row = store.latest()
    assert row["operation"] == "retrieval.vector_search"
    assert row["duration_ms"] >= 0
    assert row["attributes"]["raw_candidate_count"] == 12


# --- Store contents ---------------------------------------------------------


def test_default_status_is_success(tmp_path):
    store = LocalObservationStore(tmp_path)
    with observe_operation("rag.request", store=store):
        pass
    assert store.latest()["status"] == "success"


def test_subsystem_and_consumer_recorded(tmp_path):
    store = LocalObservationStore(tmp_path)
    with observe_operation("retrieval.vector_search", store=store, subsystem="chroma", consumer="chat"):
        pass
    row = store.latest()
    assert row["subsystem"] == "chroma"
    assert row["consumer"] == "chat"


def test_stage_durations_and_counts_recorded(tmp_path):
    store = LocalObservationStore(tmp_path)
    with observe_operation("rag.request", store=store) as op:
        op.add_stage("embedding", 5.0)
        op.add_count("candidates", 7)
    row = store.latest()
    assert row["stage_durations_ms"]["embedding"] == 5.0
    assert row["counts"]["candidates"] == 7


def test_retrieval_outcome_recorded(tmp_path):
    store = LocalObservationStore(tmp_path)
    with observe_operation("retrieval.vector_search", store=store) as op:
        op.set_retrieval(RetrievalOutcome.failure("ChromaError"))
    row = store.latest()
    assert row["retrieval"]["status"] == "error_fallback"
    assert row["retrieval"]["retrieval_error"] is True


def test_experiment_identity_recorded(tmp_path):
    store = LocalObservationStore(tmp_path)
    with observe_operation("rag.request", store=store) as op:
        op.set_experiment(ExperimentIdentity(experiment_id="exp1", pipeline_version="rag-naive-v1"))
    row = store.latest()
    assert row["experiment"]["experiment_id"] == "exp1"
    assert row["experiment"]["pipeline_version"] == "rag-naive-v1"


def test_forbidden_attribute_is_redacted_in_store(tmp_path):
    store = LocalObservationStore(tmp_path)
    with observe_operation("rag.request", store=store) as op:
        op.set("prompt", "the raw user prompt")
        op.set("raw_candidate_count", 3)
    row = store.latest()
    assert row["attributes"]["prompt"] == "[redacted]"
    assert row["attributes"]["raw_candidate_count"] == 3


# --- Exception handling -----------------------------------------------------


def test_exception_marks_error_status_and_reraises(tmp_path):
    store = LocalObservationStore(tmp_path)
    with pytest.raises(ValueError):
        with observe_operation("rag.request", store=store):
            raise ValueError("boom")
    row = store.latest()
    assert row["status"] == "error"
    assert row["attributes"].get("error.type") == "ValueError"


def test_mark_error_sets_status_without_exception(tmp_path):
    store = LocalObservationStore(tmp_path)
    with observe_operation("retrieval.vector_search", store=store) as op:
        op.mark_error("collection_missing")
    row = store.latest()
    assert row["status"] == "error"
    assert row["attributes"].get("error.type") == "collection_missing"


def test_mark_terminal_overrides_status(tmp_path):
    store = LocalObservationStore(tmp_path)
    with observe_operation("retrieval.vector_search", store=store) as op:
        op.mark_terminal("success_empty")
    assert store.latest()["status"] == "success_empty"


# --- disabled mode ----------------------------------------------------------


def test_disabled_mode_is_a_no_op_but_runs_product_code(tmp_path):
    bootstrap.initialize_observability(
        config=ObservabilityConfig(
            enabled=False,
            mode="disabled",
            otlp_endpoint="http://127.0.0.1:4318",
            signoz_ui_url="http://127.0.0.1:8080",
            artifact_root=tmp_path,
        )
    )
    ran = False
    with observe_operation("rag.request") as op:
        op.set("raw_candidate_count", 1)  # must not raise
        ran = True
    assert ran is True
    # No global store in disabled mode -> nothing written.
    assert bootstrap.get_state().local_store is None


# --- local mode global init: spans + metrics -------------------------------


def test_local_mode_creates_span_with_sanitized_attributes(tmp_path):
    bootstrap.initialize_observability(config=_local_config(tmp_path))
    with observe_operation("retrieval.vector_search") as op:
        op.set("raw_candidate_count", 9)
    spans = bootstrap.get_captured_spans()
    assert len(spans) == 1
    span = spans[0]
    assert span.name == "retrieval.vector_search"
    assert span.attributes["raw_candidate_count"] == 9
    assert span.attributes["operation"] == "retrieval.vector_search"


def test_local_mode_records_duration_histogram(tmp_path):
    bootstrap.initialize_observability(config=_local_config(tmp_path))
    with observe_operation("retrieval.vector_search", consumer="chat"):
        pass
    points = bootstrap.collect_histogram_points("operation_duration_ms")
    assert points, "expected at least one histogram data point"
    dims = points[0].attributes
    assert dims["operation"] == "retrieval.vector_search"
    # consumer is a metric-safe dimension
    assert dims.get("consumer") == "chat"
    # raw operation IDs must never become metric dimensions
    assert "trace_id" not in dims


def test_local_mode_global_store_written_without_explicit_store(tmp_path):
    bootstrap.initialize_observability(config=_local_config(tmp_path))
    with observe_operation("rag.request") as op:
        op.set("raw_candidate_count", 4)
    row = bootstrap.get_state().local_store.latest()
    assert row["operation"] == "rag.request"
    assert row["attributes"]["raw_candidate_count"] == 4


# --- correlated logging -----------------------------------------------------


def test_error_emits_correlated_log(tmp_path, caplog):
    store = LocalObservationStore(tmp_path)
    with caplog.at_level(logging.ERROR, logger="app.observability"):
        with observe_operation("retrieval.vector_search", store=store, consumer="chat") as op:
            op.mark_error("collection_missing")
    records = [r for r in caplog.records if getattr(r, "otel_operation", None) == "retrieval.vector_search"]
    assert records, "expected a correlated log record for the error"
    rec = records[0]
    assert rec.otel_status == "error"
    assert rec.otel_error_type == "collection_missing"
    assert rec.otel_consumer == "chat"


def test_success_does_not_emit_error_log(tmp_path, caplog):
    store = LocalObservationStore(tmp_path)
    with caplog.at_level(logging.ERROR, logger="app.observability"):
        with observe_operation("rag.request", store=store):
            pass
    assert not [r for r in caplog.records if getattr(r, "otel_operation", None)]


# --- span-start failure must degrade, never propagate -----------------------


def test_span_start_failure_still_runs_product_body_and_records(tmp_path, monkeypatch):
    """A tracer that fails to start/enter a span must degrade to a no-span
    operation -- the product body inside the `with` must still run untouched.
    Regression guard: span setup used to sit outside the protective try-block.
    """
    bootstrap.initialize_observability(config=_local_config(tmp_path))
    tracer = bootstrap.get_state().tracer

    def _boom(*_args, **_kwargs):
        raise RuntimeError("tracer exploded on start_as_current_span")

    monkeypatch.setattr(tracer, "start_as_current_span", _boom)

    ran = False
    result = None
    with observe_operation("retrieval.vector_search") as op:
        op.set("raw_candidate_count", 5)  # must not raise despite no span
        ran = True
        result = "product-result"

    assert ran is True
    assert result == "product-result"

    # No span was created, so no captured spans and no trace id ...
    assert bootstrap.get_captured_spans() == []
    # ... but the observation was still recorded with the product attribute.
    row = bootstrap.get_state().local_store.latest()
    assert row["operation"] == "retrieval.vector_search"
    assert row["attributes"]["raw_candidate_count"] == 5
    assert row["trace_id"] is None


# --- log export is scoped to the sanitized correlated-log path only ---------


def test_ordinary_app_logs_are_not_exported_only_sanitized_correlated_logs_are(tmp_path):
    """The OTel log handler must be scoped to ``app.observability`` so ordinary
    ``app.*`` logs (which routinely embed filesystem paths) never reach the log
    pipeline -- only ``emit_correlated_log``'s path-stripped records do.
    """
    bootstrap.initialize_observability(config=_local_config(tmp_path))
    # Drop the bootstrap's own init log so we assert on a clean pipeline.
    bootstrap.get_state().log_exporter.clear()

    secret_path = "C:\\Users\\SystemSu\\.studybuddy\\cognee\\main.py"

    # A logger OUTSIDE app.observability, exactly as cerebras_client / handlers use.
    other = logging.getLogger("app.some_other_module")
    other.setLevel(logging.INFO)
    other.error("failed reading %s during cognify", secret_path)

    assert bootstrap.get_captured_logs() == [], (
        "ordinary app.* logs must not reach the OTel log pipeline"
    )

    # The sanitized correlated-log path IS the only thing that reaches export.
    emit_correlated_log(
        "operation retrieval.vector_search finished with status=error",
        operation="retrieval.vector_search",
        status="error",
    )
    logs = bootstrap.get_captured_logs()
    assert logs, "emit_correlated_log should reach the OTel log pipeline"
    bodies = [str(getattr(r.log_record, "body", "")) for r in logs]
    assert any("retrieval.vector_search" in b for b in bodies)
    assert all(secret_path not in b for b in bodies)


# --- record_event -----------------------------------------------------------


def test_record_event_does_not_raise_inside_operation(tmp_path):
    bootstrap.initialize_observability(config=_local_config(tmp_path))
    with observe_operation("rag.request"):
        record_event("cache_hit", {"cache.state": "warm"})
    spans = bootstrap.get_captured_spans()
    assert spans[0].events, "expected the event to be attached to the span"
    assert spans[0].events[0].name == "cache_hit"