cmvr_edge_ai/tests/unit/test_framed_worker.py

from __future__ import annotations

import asyncio
import sys

import pytest

from cmvr_edge_ai.workers import (
    PROTOCOL_VERSION,
    FramedSubprocessWorker,
    FramedWorkerProtocolError,
)


_ECHO_WORKER = r"""
import json
import os
import struct
import sys
import time

length = struct.Struct("!I")

def read_frame():
    prefix = sys.stdin.buffer.read(4)
    if not prefix:
        return None, ()
    size = length.unpack(prefix)[0]
    header = json.loads(sys.stdin.buffer.read(size))
    blobs = tuple(sys.stdin.buffer.read(item) for item in header.pop("blob_lengths"))
    return header, blobs

def write_frame(header, blobs=()):
    header = dict(header)
    header["blob_lengths"] = [len(item) for item in blobs]
    encoded = json.dumps(header, separators=(",", ":")).encode()
    sys.stdout.buffer.write(length.pack(len(encoded)))
    sys.stdout.buffer.write(encoded)
    for blob in blobs:
        sys.stdout.buffer.write(blob)
    sys.stdout.buffer.flush()

write_frame({"protocol": "cmvr.external-worker/v1", "type": "ready", "pid": os.getpid()})
for header, blobs in iter(read_frame, (None, ())):
    if header["type"] == "shutdown":
        break
    if header.get("operation") == "slow":
        time.sleep(1)
    write_frame(
        {
            "protocol": "cmvr.external-worker/v1",
            "type": "response",
            "request_id": header["request_id"],
            "pid": os.getpid(),
        },
        tuple(blob[::-1] for blob in blobs),
    )
"""


_INVALID_READY_WORKER = r"""
import json
import struct
import sys
header = json.dumps({"protocol": "wrong", "type": "ready", "blob_lengths": []}).encode()
sys.stdout.buffer.write(struct.pack("!I", len(header)) + header)
sys.stdout.buffer.flush()
"""


def test_framed_worker_keeps_one_process_and_transfers_raw_blobs() -> None:
    async def exercise() -> None:
        worker = FramedSubprocessWorker(
            (sys.executable, "-u", "-c", _ECHO_WORKER),
            startup_timeout_s=2,
            shutdown_timeout_s=1,
            max_blob_bytes=1024,
        )
        try:
            ready = await worker.start()
            first = await worker.submit(
                {
                    "protocol": PROTOCOL_VERSION,
                    "type": "request",
                    "request_id": "one",
                },
                (b"image-data", b"artifact"),
                timeout_s=2,
            )
            second = await worker.submit(
                {
                    "protocol": PROTOCOL_VERSION,
                    "type": "request",
                    "request_id": "two",
                },
                (b"second",),
                timeout_s=2,
            )
            assert first.header["request_id"] == "one"
            assert first.header["pid"] == ready["pid"] == second.header["pid"]
            assert first.blobs == (b"atad-egami", b"tcafitra")
            assert second.blobs == (b"dnoces",)
            assert worker.is_alive is True
        finally:
            await worker.stop()
            await worker.stop()
        assert worker.is_alive is False

    asyncio.run(exercise())


def test_timeout_discards_uncorrelated_process_and_next_call_restarts() -> None:
    async def exercise() -> None:
        worker = FramedSubprocessWorker(
            (sys.executable, "-u", "-c", _ECHO_WORKER),
            startup_timeout_s=2,
            shutdown_timeout_s=1,
        )
        try:
            original_pid = (await worker.start())["pid"]
            with pytest.raises(asyncio.TimeoutError):
                await worker.submit(
                    {
                        "protocol": PROTOCOL_VERSION,
                        "type": "request",
                        "request_id": "slow",
                        "operation": "slow",
                    },
                    timeout_s=0.05,
                )
            assert worker.is_alive is False

            response = await worker.submit(
                {
                    "protocol": PROTOCOL_VERSION,
                    "type": "request",
                    "request_id": "after-restart",
                },
                timeout_s=2,
            )
            assert response.header["pid"] != original_pid
            assert response.header["request_id"] == "after-restart"
        finally:
            await worker.stop()

    asyncio.run(exercise())


def test_invalid_ready_handshake_is_rejected_and_child_is_reaped() -> None:
    async def exercise() -> None:
        worker = FramedSubprocessWorker(
            (sys.executable, "-u", "-c", _INVALID_READY_WORKER),
            startup_timeout_s=2,
            shutdown_timeout_s=1,
        )
        with pytest.raises(FramedWorkerProtocolError, match="protocol"):
            await worker.start()
        assert worker.is_alive is False
        await worker.stop()

    asyncio.run(exercise())


def test_framed_worker_enforces_outbound_blob_limit() -> None:
    async def exercise() -> None:
        worker = FramedSubprocessWorker(
            (sys.executable, "-u", "-c", _ECHO_WORKER),
            startup_timeout_s=2,
            shutdown_timeout_s=1,
            max_blob_bytes=3,
        )
        try:
            with pytest.raises(FramedWorkerProtocolError, match="outbound"):
                await worker.submit(
                    {
                        "protocol": PROTOCOL_VERSION,
                        "type": "request",
                        "request_id": "too-large",
                    },
                    (b"four",),
                    timeout_s=2,
                )
            # A rejected frame never reached the child, but the conservative
            # exchange boundary still resets correlation before reuse.
            assert worker.is_alive is False
        finally:
            await worker.stop()

    asyncio.run(exercise())