neuron-tai/tests/test_shard_lifecycle.py

"""DGR-022 lifecycle and structured-status contract tests."""

from __future__ import annotations

import json

import pytest

from meshnet_node.shard_lifecycle import (
    CacheExpectation,
    CacheResult,
    CapabilityResponse,
    CancellationToken,
    FlowControl,
    FlowControlLimits,
    LifecycleContractError,
    LifecycleState,
    RpcName,
    SessionFrame,
    SessionLifecycle,
    SessionPhase,
    SessionRequest,
    ShardRpcContract,
    StatusCode,
    StructuredStatus,
)


def _request(**overrides) -> SessionRequest:
    values = dict(
        schema_version=1,
        request_id="request-1",
        work_id="work-1",
        route_session="route-1",
        route_epoch=3,
        artifact_fingerprint="artifact-sha",
        runtime_fingerprint="runtime-sha",
        shard_start=0,
        shard_end=8,
        effective_start=0,
        cache_expectation=CacheExpectation.OPTIONAL,
    )
    values.update(overrides)
    return SessionRequest(**values)


def _frame(phase: SessionPhase, step: int, position: int = 0) -> SessionFrame:
    return SessionFrame(
        phase=phase,
        position=position,
        idempotency_step=step,
        payload=b"frame",
        size_bytes=5,
    )


def test_contract_roundtrip_exposes_all_lifecycle_rpcs_and_operational_hooks():
    contract = ShardRpcContract()
    restored = ShardRpcContract.from_bytes(contract.to_bytes())

    assert restored == contract
    assert restored.methods == tuple(RpcName)
    assert restored.deadlines.session_idle_seconds > restored.deadlines.health_seconds
    assert restored.flow_control.max_inflight_frames == 32
    assert restored.tls_auth.tls_required is True


def test_supported_version_roundtrip_and_status_are_stable():
    status = StructuredStatus(StatusCode.RESOURCE_EXHAUSTED, "window full", retryable=True, details={"limit": "32"})
    restored = StructuredStatus.from_dict(status.to_dict())

    assert restored == status
    assert CapabilityResponse(
        status=StructuredStatus(StatusCode.OK, "ready"),
        schema_version=1,
        supported_versions=(1,),
    ).validate() is None


@pytest.mark.parametrize("version", [0, 2, 99])
def test_unsupported_versions_fail_closed(version):
    with pytest.raises(LifecycleContractError) as exc:
        ShardRpcContract.from_bytes(json.dumps({
            "schema": "meshnet.shard-runtime",
            "version": version,
            "methods": [method.value for method in RpcName],
        }).encode())
    assert exc.value.status.code is StatusCode.UNSUPPORTED_VERSION

    with pytest.raises(LifecycleContractError) as exc:
        _request(schema_version=version).validate()
    assert exc.value.status.code is StatusCode.UNSUPPORTED_VERSION


def test_malformed_lifecycle_transitions_fail_closed():
    session = SessionLifecycle(_request())

    with pytest.raises(LifecycleContractError) as exc:
        session.apply(_frame(SessionPhase.DECODE, 0))
    assert exc.value.status.code is StatusCode.MALFORMED_LIFECYCLE
    assert session.state is LifecycleState.OPEN

    session.apply(_frame(SessionPhase.PREFILL, 0))
    session.apply(_frame(SessionPhase.DECODE, 1, position=1))
    assert session.state is LifecycleState.DECODING

    with pytest.raises(LifecycleContractError) as exc:
        session.apply(_frame(SessionPhase.PREFILL, 2))
    assert exc.value.status.code is StatusCode.MALFORMED_LIFECYCLE


def test_duplicate_and_non_monotonic_idempotency_steps_are_rejected():
    session = SessionLifecycle(_request())
    session.apply(_frame(SessionPhase.PREFILL, 4))

    with pytest.raises(LifecycleContractError) as exc:
        session.apply(_frame(SessionPhase.PREFILL, 4))
    assert exc.value.status.code is StatusCode.ALREADY_EXISTS

    with pytest.raises(LifecycleContractError) as exc:
        session.apply(_frame(SessionPhase.PREFILL, 3))
    assert exc.value.status.code is StatusCode.MALFORMED_LIFECYCLE


def test_cancel_propagates_to_waiters_and_release_is_idempotent():
    session = SessionLifecycle(_request())
    assert session.cancel("client disconnected").code is StatusCode.CANCELLED
    assert session.cancellation.cancelled is True
    assert session.state is LifecycleState.CANCELLED

    with pytest.raises(LifecycleContractError) as exc:
        session.apply(_frame(SessionPhase.DECODE, 0))
    assert exc.value.status.code is StatusCode.CANCELLED

    assert session.release().code is StatusCode.OK
    assert session.state is LifecycleState.RELEASED
    assert session.release().code is StatusCode.OK


def test_bounded_flow_control_rejects_oversize_and_full_windows():
    flow = FlowControl(FlowControlLimits(max_inflight_frames=1, max_inflight_bytes=10, max_frame_bytes=8))
    flow.acquire(8)

    with pytest.raises(LifecycleContractError) as exc:
        flow.acquire(1)
    assert exc.value.status.code is StatusCode.RESOURCE_EXHAUSTED
    assert exc.value.status.retryable is True

    with pytest.raises(LifecycleContractError) as exc:
        flow.acquire(9)
    assert exc.value.status.code is StatusCode.RESOURCE_EXHAUSTED

    flow.release(8)
    assert flow.outstanding == (0, 0)


def test_flow_control_wait_honours_cancellation():
    token = CancellationToken()
    flow = FlowControl(FlowControlLimits(max_inflight_frames=1, max_inflight_bytes=10, max_frame_bytes=8))
    flow.acquire(4)
    token.cancel()

    with pytest.raises(LifecycleContractError) as exc:
        flow.acquire(1, wait=True, cancelled=lambda: token.cancelled)
    assert exc.value.status.code is StatusCode.CANCELLED


def test_cache_expectation_and_result_are_explicit():
    request = _request(cache_expectation=CacheExpectation.REQUIRED)
    assert request.cache_expectation is CacheExpectation.REQUIRED
    assert CacheResult.MISS.value == "MISS"


def test_invalid_contract_methods_and_tls_hooks_fail_closed():
    with pytest.raises(LifecycleContractError) as exc:
        ShardRpcContract(methods=(RpcName.HEALTH,)).validate()
    assert exc.value.status.code is StatusCode.FAILED_PRECONDITION

    with pytest.raises(ValueError, match="server_name"):
        ShardRpcContract(
            tls_auth=type(ShardRpcContract().tls_auth)(tls_required=True, server_name="")
        ).validate()