NODEDC_MISSION_CORE/tests/test_pipeline_telemetry.py

from __future__ import annotations

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

import pytest

from k1link.compute.pipeline_telemetry import (
    MAX_PAYLOAD_BYTES,
    JsonlPipelineTelemetrySink,
    MqttPipelineTelemetrySink,
    PipelineTelemetryEmitter,
    PipelineTelemetryError,
    PipelineTelemetryIdentity,
    build_pipeline_run_telemetry_document,
    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,
        activation_count=7,
        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["payload"]["event"]["activation_count"] == 7
    assert stored["payload"]["stage_metrics"]["predict"]["activations"] == 7
    assert stored["authority"]["commands_enabled"] is False


def test_telegraf_tail_wrapper_restores_the_native_pipeline_document() -> None:
    identity = _identity()
    native = build_pipeline_telemetry_document(
        identity=identity,
        stage_id="tracking",
        state="completed",
        duration_ms=4.25,
        observed_at_utc="2026-07-28T12:00:00Z",
    )
    record = {
        "schema_version": "missioncore.pipeline-telemetry-record/v1",
        "topic": identity.topic,
        "payload": native,
    }
    telegraf = {
        "name": "missioncore_pipeline_event",
        "timestamp": 1785240000,
        "tags": {
            "agent_id": identity.agent_id,
            "contour_id": identity.contour_id,
            "node_id": identity.node_id,
        },
        "fields": {
            "value": json.dumps(record, separators=(",", ":")),
        },
    }

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

    assert row[5] == "pipeline"
    assert row[7] == "missioncore.agent-pipeline-telemetry/v1"
    assert row[9:13] == ("E41", "run-001", "request-001", 17)
    assert json.loads(row[13]) == native

    record["topic"] = "mission-core/v1/contours/other/agents/other/pipeline"
    telegraf["fields"]["value"] = json.dumps(record)
    with pytest.raises(ValueError, match="pipeline event record"):
        _normalizer()._normalize(identity.topic, json.dumps(telegraf).encode())


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_run_events_record_explicit_terminal_outcome() -> None:
    identity = _identity()
    started = build_pipeline_run_telemetry_document(
        identity=identity,
        state="started",
        observed_at_utc="2026-07-28T12:00:00Z",
    )
    completed = build_pipeline_run_telemetry_document(
        identity=identity,
        state="completed",
        duration_ms=12_345.5,
        exit_code=2,
        observed_at_utc="2026-07-28T12:00:12Z",
    )

    assert started["payload"]["state"] == "busy"
    assert started["payload"]["active_request_id"] == "request-001"
    assert completed["event_type"] == "run"
    assert completed["run_state"] == "completed"
    assert completed["payload"]["active_request_id"] is None
    assert completed["payload"]["event"] == {
        "event_type": "run",
        "state": "completed",
        "duration_ms": 12_345.5,
        "exit_code": 2,
        "error_type": None,
    }

    with pytest.raises(PipelineTelemetryError, match="completed run"):
        build_pipeline_run_telemetry_document(
            identity=identity,
            state="completed",
            duration_ms=1.0,
        )


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_jsonl_sink_rotates_to_content_addressed_segment(tmp_path: Path) -> None:
    path = tmp_path / "pipeline-telemetry.jsonl"
    maximum = MAX_PAYLOAD_BYTES + 4096
    existing = (b"{}\n" * (maximum // 3))[: maximum - 32]
    path.write_bytes(existing)
    sink = JsonlPipelineTelemetrySink(path, max_bytes=maximum, max_segments=2)

    sink.publish(_identity().topic, b"{}")

    segments = tuple(tmp_path.glob("pipeline-telemetry.*.jsonl"))
    assert len(segments) == 1
    assert segments[0].read_bytes() == existing
    assert segments[0].stem.split(".")[-1] == hashlib.sha256(existing).hexdigest()
    assert json.loads(path.read_text())["payload"] == {}


def test_jsonl_sink_refuses_to_delete_unacknowledged_segment_at_bound(
    tmp_path: Path,
) -> None:
    path = tmp_path / "pipeline-telemetry.jsonl"
    maximum = MAX_PAYLOAD_BYTES + 4096
    active = b"x" * maximum
    path.write_bytes(active)
    retained = tmp_path / f"pipeline-telemetry.{'a' * 64}.jsonl"
    retained.write_bytes(b"retained")
    sink = JsonlPipelineTelemetrySink(path, max_bytes=maximum, max_segments=1)

    with pytest.raises(PipelineTelemetryError, match="segment bound reached"):
        sink.publish(_identity().topic, b"{}")

    assert path.read_bytes() == active
    assert retained.read_bytes() == b"retained"


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",
        )