neuron-tai/tests/test_native_activation_seam.py

"""DGR-042 seam tests with a deterministic fake generated worker."""

from __future__ import annotations

from collections.abc import Iterator
import threading

import pytest

from meshnet_node.native_activation_seam import (
    NATIVE_RELAY_PATH,
    NativeActivationBufferFull,
    NativeActivationDisconnected,
    NativeActivationSeam,
    NativeFrameContext,
)
from meshnet_node.native_protocol import pb


def _context(**changes: object) -> NativeFrameContext:
    values: dict[str, object] = dict(
        request_id="billing-request-7", node_id="node-tail", route_session_id="route-9",
        route_epoch=4, work_id="work-3", deadline_unix_nanos=987654321,
    )
    values.update(changes)
    return NativeFrameContext(**values)


def _open() -> pb.SessionRequest:
    return pb.SessionRequest(open=pb.SessionOpen(
        schema_version=pb.SCHEMA_VERSION_1, route_session_id="route-9", route_epoch=4,
    ))


def _chunk() -> pb.SessionRequest:
    return pb.SessionRequest(chunk=pb.ActivationChunk(envelope=pb.Envelope(
        schema_version=pb.SCHEMA_VERSION_1, route_session_id="route-9", route_epoch=4,
        work_id="work-3", deadline_unix_nanos=987654321,
    )))


def _ack(request: pb.SessionRequest) -> pb.SessionResponse:
    if request.WhichOneof("kind") == "open":
        return pb.SessionResponse(accepted=pb.SessionAccepted(
            schema_version=pb.SCHEMA_VERSION_1, route_session_id="route-9", route_epoch=4,
        ))
    route, epoch, work, _ = ("route-9", 4, "work-3", 0)
    del route, epoch
    return pb.SessionResponse(ack=pb.Ack(work_id=work, idempotency_step=1))


class _FakeGrpcWorker:
    def __init__(self, *, block: bool = False) -> None:
        self.calls = 0
        self.received: list[pb.SessionRequest] = []
        self.started = threading.Event()
        self.consumed = threading.Event()
        self.release = threading.Event()
        self.block = block

    def Session(self, requests: Iterator[pb.SessionRequest]):
        self.calls += 1
        self.started.set()
        for request in requests:
            self.received.append(request)
            self.consumed.set()
            if self.block:
                self.release.wait(1)
            yield _ack(request)


def test_direct_uses_one_long_lived_grpc_stream_and_preserves_correlation():
    worker = _FakeGrpcWorker()
    telemetry = []
    seam = NativeActivationSeam(_context(), direct_stub=worker, telemetry=telemetry.append)
    try:
        seam.send(_open())
        seam.send(_chunk())
        assert seam.receive(1).WhichOneof("kind") == "accepted"
        assert seam.receive(1).ack.work_id == "work-3"
        assert worker.calls == 1
        assert [frame.SerializeToString() for frame in worker.received] == [
            _open().SerializeToString(), _chunk().SerializeToString()
        ]
        assert telemetry[-1].request_id == "billing-request-7"
        assert telemetry[-1].node_id == "node-tail"
    finally:
        seam.close()


def test_relay_carries_byte_identical_protobuf_frames_and_all_correlation_headers():
    captured: list[tuple[str, bytes, dict[str, str]]] = []

    def relay(path: str, body: bytes, headers: dict[str, str]):
        captured.append((path, body, headers))
        request = pb.SessionRequest()
        request.ParseFromString(body)
        return 200, {}, _ack(request).SerializeToString()

    seam = NativeActivationSeam(_context(), relay_request=relay)
    response = seam.send(_chunk())
    assert response is not None and response.ack.work_id == "work-3"
    path, body, headers = captured[0]
    assert path == NATIVE_RELAY_PATH
    assert body == _chunk().SerializeToString()
    assert headers == {
        "Content-Type": "application/x-protobuf", "X-Meshnet-Native-Frame": "shard-runtime/v1",
        "X-Meshnet-Request-Id": "billing-request-7", "X-Meshnet-Node-Id": "node-tail",
        "X-Meshnet-Session": "route-9", "X-Meshnet-Route-Epoch": "4",
        "X-Meshnet-Work-Id": "work-3", "X-Meshnet-Deadline-Unix-Nanos": "987654321",
        "X-Meshnet-Activation-Id": "billing-request-7",
    }


def test_relay_disconnect_is_uncertain_and_is_never_replayed():
    calls = 0

    def disconnected(*_):
        nonlocal calls
        calls += 1
        raise OSError("relay vanished")

    seam = NativeActivationSeam(_context(), relay_request=disconnected)
    with pytest.raises(NativeActivationDisconnected, match="uncertain"):
        seam.send(_chunk())
    with pytest.raises(NativeActivationDisconnected):
        seam.send(_chunk())
    assert calls == 1


def test_cancellation_uses_the_same_opaque_relay_contract():
    received = []

    def relay(_path, body, _headers):
        request = pb.SessionRequest()
        request.ParseFromString(body)
        received.append(request)
        return 200, {}, pb.SessionResponse(ack=pb.Ack(work_id="work-3")).SerializeToString()

    seam = NativeActivationSeam(_context(), relay_request=relay)
    response = seam.cancel("client disconnected")
    assert response is not None
    assert received[0].cancel.work_id == "work-3"
    assert received[0].cancel.reason == "client disconnected"


def test_direct_request_buffer_is_bounded():
    worker = _FakeGrpcWorker(block=True)
    seam = NativeActivationSeam(_context(), direct_stub=worker, max_buffered_frames=1)
    try:
        assert worker.started.wait(1)
        seam.send(_open())
        assert worker.consumed.wait(1)
        seam.send(_chunk())
        with pytest.raises(NativeActivationBufferFull):
            seam.send(_chunk())
    finally:
        worker.release.set()
        seam.close()