cmvr_edge_ai/tests/integration/test_runtime.py

from __future__ import annotations

import asyncio
from typing import Any

from cmvr_edge_ai.core import (
    ComponentContext,
    Edge,
    Envelope,
    PipelineRuntime,
    RuntimeState,
    Sink,
    Source,
)


class ListSource(Source):
    def __init__(self, items: list[str]) -> None:
        self.items = items

    async def messages(self):  # type: ignore[no-untyped-def]
        for sequence, item in enumerate(self.items):
            yield Envelope(item, source_id="source", sequence=sequence)


class CollectSink(Sink):
    def __init__(self, delay_s: float = 0.0) -> None:
        self.delay_s = delay_s
        self.envelopes: list[Envelope[Any]] = []
        self.input_ports: list[str] = []
        self.stop_calls = 0

    async def consume(self, envelope: Envelope[Any], input_port: str = "input") -> None:
        if self.delay_s:
            await asyncio.sleep(self.delay_s)
        self.envelopes.append(envelope)
        self.input_ports.append(input_port)

    async def stop(self) -> None:
        self.stop_calls += 1


class BurstThenWaitSource(Source):
    def __init__(self, items: list[str]) -> None:
        self.items = items
        self.produced = asyncio.Event()
        self.release = asyncio.Event()
        self.stop_calls = 0

    async def setup(self, context: ComponentContext) -> None:
        self.shutdown_event = context.shutdown_event

    async def messages(self):  # type: ignore[no-untyped-def]
        for sequence, item in enumerate(self.items):
            yield Envelope(item, source_id="burst", sequence=sequence)
        self.produced.set()
        await self.release.wait()

    async def stop(self) -> None:
        self.stop_calls += 1


class SetupBlockingSource(Source):
    def __init__(
        self,
        setup_entered: asyncio.Event,
        allow_setup: asyncio.Event,
        run_forever: asyncio.Event,
    ) -> None:
        self.setup_entered = setup_entered
        self.allow_setup = allow_setup
        self.run_forever = run_forever
        self.stop_calls = 0

    async def setup(self, context: ComponentContext) -> None:
        del context
        self.setup_entered.set()
        await self.allow_setup.wait()

    async def messages(self):  # type: ignore[no-untyped-def]
        await self.run_forever.wait()
        if False:
            yield Envelope("unreachable")

    async def stop(self) -> None:
        self.stop_calls += 1


class BlockingStopSource(Source):
    def __init__(
        self,
        run_forever: asyncio.Event,
        stop_entered: asyncio.Event,
        stop_cancelled: asyncio.Event,
    ) -> None:
        self.run_forever = run_forever
        self.stop_entered = stop_entered
        self.stop_cancelled = stop_cancelled

    async def messages(self):  # type: ignore[no-untyped-def]
        await self.run_forever.wait()
        if False:
            yield Envelope("unreachable")

    async def stop(self) -> None:
        self.stop_entered.set()
        try:
            await asyncio.Event().wait()
        except asyncio.CancelledError:
            self.stop_cancelled.set()
            raise


def test_runtime_fan_out_delivers_each_envelope_to_each_edge() -> None:
    async def exercise():  # type: ignore[no-untyped-def]
        source = ListSource(["one", "two", "three"])
        left = CollectSink()
        right = CollectSink()
        runtime = PipelineRuntime(
            {"source": source, "left": left, "right": right},
            [
                Edge("source", "left", capacity=2, name="left-edge"),
                Edge("source", "right", capacity=2, name="right-edge"),
            ],
            pipeline_id="fanout",
        )
        await runtime.run()
        return runtime, left, right

    runtime, left, right = asyncio.run(exercise())

    assert runtime.state is RuntimeState.STOPPED
    assert [item.payload for item in left.envelopes] == ["one", "two", "three"]
    assert [item.payload for item in right.envelopes] == ["one", "two", "three"]
    assert all(
        left_item is right_item
        for left_item, right_item in zip(left.envelopes, right.envelopes)
    )
    assert left.input_ports == ["input"] * 3
    assert set(runtime.edge_stats()) == {"left-edge", "right-edge"}


def test_graceful_stop_drains_already_accepted_messages_and_stops_once() -> None:
    async def exercise():  # type: ignore[no-untyped-def]
        source = BurstThenWaitSource(["one", "two", "three"])
        sink = CollectSink(delay_s=0.01)
        runtime = PipelineRuntime(
            {"source": source, "sink": sink},
            [Edge("source", "sink", capacity=3, name="drain-edge")],
            pipeline_id="graceful",
            shutdown_timeout_s=1,
        )

        await runtime.start()
        await asyncio.wait_for(source.produced.wait(), timeout=1)
        await runtime.stop(graceful=True)
        return source, sink, runtime

    source, sink, runtime = asyncio.run(exercise())

    assert [item.payload for item in sink.envelopes] == ["one", "two", "three"]
    assert source.stop_calls == 1
    assert sink.stop_calls == 1
    assert runtime.state is RuntimeState.STOPPED
    stats = runtime.edge_stats()["drain-edge"]
    assert stats.enqueued == stats.dequeued == 3
    assert stats.discarded_on_close == 0
    assert stats.closed is True


def test_stop_racing_start_waits_for_startup_then_tears_down_once() -> None:
    async def exercise():  # type: ignore[no-untyped-def]
        setup_entered = asyncio.Event()
        allow_setup = asyncio.Event()
        source = SetupBlockingSource(
            setup_entered,
            allow_setup,
            asyncio.Event(),
        )
        sink = CollectSink()
        runtime = PipelineRuntime(
            {"source": source, "sink": sink},
            [Edge("source", "sink")],
            pipeline_id="start-stop-race",
        )

        start_task = asyncio.create_task(runtime.start())
        await asyncio.wait_for(setup_entered.wait(), timeout=1)
        stop_task = asyncio.create_task(runtime.stop(graceful=False, timeout_s=1))
        await asyncio.sleep(0)
        assert runtime.state is RuntimeState.STARTING
        assert stop_task.done() is False

        allow_setup.set()
        await asyncio.wait_for(asyncio.gather(start_task, stop_task), timeout=1)
        return runtime, source, sink

    runtime, source, sink = asyncio.run(exercise())

    assert runtime.state is RuntimeState.STOPPED
    assert runtime.shutdown_event.is_set()
    assert source.stop_calls == 1
    assert sink.stop_calls == 1


def test_blocking_component_stop_is_bounded_by_shutdown_timeout() -> None:
    async def exercise():  # type: ignore[no-untyped-def]
        stop_entered = asyncio.Event()
        stop_cancelled = asyncio.Event()
        source = BlockingStopSource(
            asyncio.Event(),
            stop_entered,
            stop_cancelled,
        )
        sink = CollectSink()
        runtime = PipelineRuntime(
            {"source": source, "sink": sink},
            [Edge("source", "sink")],
            pipeline_id="bounded-stop",
        )
        await runtime.start()

        loop = asyncio.get_running_loop()
        started_at = loop.time()
        await runtime.stop(graceful=False, timeout_s=0.05)
        elapsed = loop.time() - started_at
        await asyncio.sleep(0)
        return runtime, stop_entered, stop_cancelled, elapsed

    runtime, stop_entered, stop_cancelled, elapsed = asyncio.run(exercise())

    assert stop_entered.is_set()
    assert stop_cancelled.is_set()
    assert elapsed < 0.5
    assert runtime.state is RuntimeState.STOPPED