neuron-tai/tests/test_native_worker_supervisor.py

"""Model-free DGR-040 supervision tests using a deterministic fake worker."""

from __future__ import annotations

import hashlib
import sys
import time
from pathlib import Path

import pytest

from meshnet_node.native_worker_supervisor import (
    NativeWorkerError,
    NativeWorkerProbe,
    NativeWorkerSpec,
    NativeWorkerSupervisor,
)


def _write_fake_worker(path: Path) -> None:
    path.write_text(
        """import os
import signal
import sys
import time

stop = False
def terminate(*_):
    global stop
    stop = True
signal.signal(signal.SIGTERM, terminate)
print('ShardRuntime worker listening on ' + os.environ['MESHNET_SHARD_LISTEN_ADDR'], flush=True)
if os.environ.get('MESHNET_INJECT_PROCESS_DEATH_AFTER_EXECUTIONS'):
    marker = os.environ.get('MESHNET_FAKE_CRASH_ONCE_FILE')
    if not marker or not os.path.exists(marker):
        if marker:
            open(marker, 'w').close()
        time.sleep(0.05)
        print('deterministic injected worker death', file=sys.stderr, flush=True)
        raise SystemExit(70)
while not stop:
    time.sleep(0.01)
print('ShardRuntime worker shut down cleanly', flush=True)
""",
        encoding="utf-8",
    )


def _spec(tmp_path: Path, **changes: object) -> NativeWorkerSpec:
    artifact = tmp_path / "fixture.gguf"
    artifact.write_bytes(b"fixture artifact")
    fake = tmp_path / "fake_worker.py"
    _write_fake_worker(fake)
    values: dict[str, object] = {
        "binary": Path(sys.executable),
        "binary_digest": hashlib.sha256(Path(sys.executable).read_bytes()).hexdigest(),
        "args": (str(fake),),
        "listen_address": "fake-worker:12345",
        "artifact_path": artifact,
        "artifact_digest": hashlib.sha256(artifact.read_bytes()).hexdigest(),
        "recipe_digest": "a" * 64,
        "recipe_id": "fixture",
        "recipe_version": "1",
        "catalogue_version": "test",
        "shard_start": 2,
        "shard_end": 5,
    }
    values.update(changes)
    return NativeWorkerSpec(**values)  # type: ignore[arg-type]


def _probe(spec: NativeWorkerSpec, _timeout: float) -> NativeWorkerProbe:
    return NativeWorkerProbe(
        artifact_digest=spec.artifact_digest,
        recipe_digest=spec.recipe_digest,
        recipe_id=spec.recipe_id,
        recipe_version=spec.recipe_version,
        catalogue_version=spec.catalogue_version,
        shard_start=spec.shard_start,
        shard_end=spec.shard_end,
        serving=True,
    )


def _eventually(predicate, timeout: float = 2.0) -> bool:
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        if predicate():
            return True
        time.sleep(0.01)
    return predicate()


def test_start_verifies_identity_captures_logs_and_stops_gracefully(tmp_path):
    events: list[tuple[str, str]] = []
    supervisor = NativeWorkerSupervisor(
        _spec(tmp_path),
        probe=_probe,
        readiness_timeout=1,
        shutdown_timeout=1,
        kill_timeout=1,
        on_available=lambda reason: events.append(("available", reason)),
        on_unavailable=lambda reason: events.append(("unavailable", reason)),
    )

    result = supervisor.start()
    assert result.serving and supervisor.available
    assert events == [("available", "worker ready and identity verified")]
    assert any("ShardRuntime worker listening" in line for line in supervisor.logs)

    supervisor.stop()
    assert not supervisor.available
    assert events[-1] == ("unavailable", "worker stopped")
    assert _eventually(lambda: any("shut down cleanly" in line for line in supervisor.logs))


def test_start_refuses_changed_artifact_before_spawning(tmp_path):
    spec = _spec(tmp_path, artifact_digest="0" * 64)
    supervisor = NativeWorkerSupervisor(spec, probe=_probe, readiness_timeout=1)

    with pytest.raises(NativeWorkerError, match="artifact digest"):
        supervisor.start()
    assert supervisor.pid is None


def test_start_refuses_changed_binary_before_spawning(tmp_path):
    spec = _spec(tmp_path, binary_digest="0" * 64)
    supervisor = NativeWorkerSupervisor(spec, probe=_probe, readiness_timeout=1)

    with pytest.raises(NativeWorkerError, match="binary digest"):
        supervisor.start()
    assert supervisor.pid is None


def test_probe_identity_mismatch_never_makes_capability_available(tmp_path):
    spec = _spec(tmp_path)

    def wrong_probe(actual: NativeWorkerSpec, timeout: float) -> NativeWorkerProbe:
        result = _probe(actual, timeout)
        return NativeWorkerProbe(**{**result.__dict__, "shard_end": actual.shard_end + 1})

    supervisor = NativeWorkerSupervisor(spec, probe=wrong_probe, readiness_timeout=1, shutdown_timeout=1)
    with pytest.raises(NativeWorkerError, match="identity/range"):
        supervisor.start()
    assert not supervisor.available


def test_deterministic_worker_death_withdraws_then_restart_recovers(tmp_path):
    unavailable: list[str] = []
    spec = _spec(
        tmp_path,
        extra_environment={
            "MESHNET_INJECT_PROCESS_DEATH_AFTER_EXECUTIONS": "1",
            "MESHNET_FAKE_CRASH_ONCE_FILE": str(tmp_path / "crashed-once"),
        },
    )
    supervisor = NativeWorkerSupervisor(
        spec,
        probe=_probe,
        readiness_timeout=1,
        health_interval=0.01,
        shutdown_timeout=1,
        kill_timeout=1,
        on_unavailable=unavailable.append,
    )
    supervisor.start()
    assert _eventually(lambda: not supervisor.available)
    assert "code 70" in supervisor.unavailable_reason
    assert any("deterministic injected worker death" in line for line in supervisor.logs)

    supervisor.restart()
    assert supervisor.available
    supervisor.stop()


def test_health_loss_withdraws_only_native_capability(tmp_path):
    healthy = True
    unavailable: list[str] = []
    spec = _spec(tmp_path)

    def health_probe(actual: NativeWorkerSpec, timeout: float) -> NativeWorkerProbe:
        result = _probe(actual, timeout)
        return NativeWorkerProbe(**{**result.__dict__, "serving": healthy})

    supervisor = NativeWorkerSupervisor(
        spec,
        probe=health_probe,
        readiness_timeout=1,
        on_unavailable=unavailable.append,
    )
    supervisor.start()
    healthy = False
    assert not supervisor.check_health()
    assert not supervisor.available
    assert unavailable and "health lost" in unavailable[-1]
    supervisor.stop()