NODEDC_MISSION_CORE/tests/test_pipeline_telemetry.py

from __future__ import annotations

import importlib.util
import json
from pathlib import Path
from types import ModuleType, SimpleNamespace

import pytest

from k1link.compute.pipeline_telemetry import (
    JsonlPipelineTelemetrySink,
    MqttPipelineTelemetrySink,
    PipelineTelemetryEmitter,
    PipelineTelemetryError,
    PipelineTelemetryIdentity,
    build_pipeline_telemetry_document,
)

REPOSITORY_ROOT = Path(__file__).resolve().parents[1]
NORMALIZER_PATH = (
    REPOSITORY_ROOT
    / "deploy"
    / "telemetry-plane"
    / "normalizer"
    / "normalizer.py"
)


def _normalizer() -> ModuleType:
    spec = importlib.util.spec_from_file_location(
        "missioncore_pipeline_telemetry_normalizer",
        NORMALIZER_PATH,
    )
    assert spec is not None and spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def _identity() -> PipelineTelemetryIdentity:
    return PipelineTelemetryIdentity(
        contour_id="worker-006",
        agent_id="mission-core-worker",
        node_id="DESKTOP-OPJ8J04",
        lab_id="E41",
        run_id="run-001",
        request_id="request-001",
        source_id="ravnoves00",
        source_package_id="e41-predictor-package-example",
        method_id="frozen-e40-predictor/v1",
        frame_index=17,
    )


def test_pipeline_document_is_accepted_without_losing_stage_identity() -> None:
    identity = _identity()
    document = build_pipeline_telemetry_document(
        identity=identity,
        stage_id="predict",
        state="completed",
        duration_ms=125.5,
        input_count=89,
        output_count=89,
        queue_wait_ms=2.25,
        observed_at_utc="2026-07-28T12:00:00Z",
    )

    row = _normalizer()._normalize(
        identity.topic,
        json.dumps(document).encode(),
    )

    assert row[1:5] == (
        "worker-006",
        "mission-core-worker",
        "DESKTOP-OPJ8J04",
        "pipeline",
    )
    assert row[9:13] == ("E41", "run-001", "request-001", 17)
    assert json.loads(row[6]) == {
        "lab_id": "E41",
        "method_id": "frozen-e40-predictor/v1",
        "request_id": "request-001",
        "run_id": "run-001",
        "source_id": "ravnoves00",
        "source_package_id": "e41-predictor-package-example",
        "stage_id": "predict",
        "stage_state": "completed",
    }
    stored = json.loads(row[13])
    assert stored["payload"]["event"]["duration_ms"] == 125.5
    assert stored["authority"]["commands_enabled"] is False


def test_stage_context_emits_terminal_event_and_preserves_failure() -> None:
    published: list[tuple[str, dict[str, object]]] = []

    class Sink:
        def publish(self, topic: str, payload: bytes) -> None:
            published.append((topic, json.loads(payload)))

    ticks = iter((1_000_000_000, 1_125_500_000))
    emitter = PipelineTelemetryEmitter(
        identity=_identity(),
        sink=Sink(),
        clock_ns=lambda: next(ticks),
    )
    with emitter.stage("predict", input_count=89) as outcome:
        outcome.output_count = 89

    assert [document["stage_state"] for _, document in published] == [
        "started",
        "completed",
    ]
    assert published[1][1]["payload"]["event"]["duration_ms"] == 125.5
    assert published[1][1]["payload"]["event"]["output_count"] == 89

    failure_ticks = iter((2_000_000_000, 2_001_000_000))
    failure_emitter = PipelineTelemetryEmitter(
        identity=_identity(),
        sink=Sink(),
        clock_ns=lambda: next(failure_ticks),
    )
    with (
        pytest.raises(ValueError, match="source failure"),
        failure_emitter.stage("evaluate"),
    ):
        raise ValueError("source failure")
    assert published[-1][1]["stage_state"] == "failed"
    assert published[-1][1]["payload"]["event"]["error_type"] == "ValueError"
    assert "source failure" not in json.dumps(published[-1][1])


def test_jsonl_sink_records_topic_bound_documents(tmp_path: Path) -> None:
    path = tmp_path / "telemetry" / "e41.jsonl"
    identity = _identity()
    sink = JsonlPipelineTelemetrySink(path)
    document = build_pipeline_telemetry_document(
        identity=identity,
        stage_id="package",
        state="completed",
        duration_ms=1.0,
    )

    sink.publish(identity.topic, json.dumps(document).encode())

    record = json.loads(path.read_text(encoding="utf-8"))
    assert record["schema_version"] == "missioncore.pipeline-telemetry-record/v1"
    assert record["topic"] == identity.topic
    assert record["payload"]["stage_id"] == "package"
    assert path.stat().st_mode & 0o077 == 0


def test_mqtt_sink_uses_qos_one_without_retention() -> None:
    calls: list[tuple[str, bytes, int, bool]] = []

    class Client:
        def publish(
            self,
            topic: str,
            payload: bytes,
            qos: int,
            retain: bool,
        ) -> SimpleNamespace:
            calls.append((topic, payload, qos, retain))
            return SimpleNamespace(rc=0)

    MqttPipelineTelemetrySink(Client()).publish("topic", b"payload")

    assert calls == [("topic", b"payload", 1, False)]


def test_pipeline_telemetry_rejects_unsafe_identity_and_invalid_metrics() -> None:
    with pytest.raises(PipelineTelemetryError, match="contour_id"):
        PipelineTelemetryIdentity(
            contour_id="../worker",
            agent_id="agent",
            node_id="node",
            lab_id="E41",
            run_id="run",
            source_id="source",
            source_package_id="package",
            method_id="method",
        )
    with pytest.raises(PipelineTelemetryError, match="duration"):
        build_pipeline_telemetry_document(
            identity=_identity(),
            stage_id="predict",
            state="completed",
        )