beaver_project/app-instance/backend/tests/unit/test_gateway_channels.py

import asyncio
from dataclasses import dataclass, field
from typing import Any

from beaver.foundation.events import InboundMessage, MessageBus
from beaver.interfaces.channels import ChannelManager, MemoryChannelAdapter
from beaver.interfaces.gateway.main import run_gateway
from beaver.services.agent_service import AgentService


@dataclass(slots=True)
class FakeResult:
    session_id: str
    run_id: str = "run-1"
    output_text: str = ""
    finish_reason: str = "stop"
    provider_name: str | None = "fake"
    model: str | None = "fake-model"
    usage: dict[str, Any] = field(default_factory=dict)


class FakeService:
    is_running = True

    async def submit_direct(self, message: str, **kwargs: Any) -> FakeResult:
        return FakeResult(
            session_id=kwargs.get("session_id") or "s1",
            output_text=f"echo:{message}",
        )

    async def handle_inbound_message(self, inbound: InboundMessage):
        result = await self.submit_direct(inbound.content, session_id=inbound.session_id)
        return AgentService.build_outbound_message(inbound, result)


class SlowService:
    is_running = True

    async def submit_direct(self, message: str, **kwargs: Any) -> FakeResult:
        await asyncio.sleep(10)
        return FakeResult(session_id=kwargs.get("session_id") or "s1")

    async def handle_inbound_message(self, inbound: InboundMessage):
        result = await self.submit_direct(inbound.content, session_id=inbound.session_id)
        return AgentService.build_outbound_message(inbound, result)


class InvalidService:
    is_running = True


def test_gateway_routes_memory_channel_roundtrip() -> None:
    async def run() -> None:
        bus = MessageBus()
        channel = MemoryChannelAdapter(bus)
        stop_event = asyncio.Event()
        task = asyncio.create_task(
            run_gateway(
                service=FakeService(),
                manage_service_lifecycle=False,
                bus=bus,
                channels=[channel],
                stop_event=stop_event,
            )
        )

        await channel.publish_text("hello", session_id="s1")
        for _ in range(40):
            if channel.sent_messages:
                break
            await asyncio.sleep(0.05)

        assert channel.sent_messages
        message = channel.sent_messages[0]
        assert message.content == "echo:hello"
        assert message.session_id == "s1"
        assert message.finish_reason == "stop"

        stop_event.set()
        await asyncio.wait_for(task, timeout=2)

    asyncio.run(run())


def test_gateway_delivers_cancelled_outbound_to_channel() -> None:
    async def run() -> None:
        bus = MessageBus()
        channel = MemoryChannelAdapter(bus)
        stop_event = asyncio.Event()
        task = asyncio.create_task(
            run_gateway(
                service=SlowService(),
                manage_service_lifecycle=False,
                bus=bus,
                channels=[channel],
                stop_event=stop_event,
            )
        )

        await channel.publish_text("slow", session_id="s1")
        await asyncio.sleep(0.05)
        stop_event.set()
        await asyncio.wait_for(task, timeout=3)

        assert channel.sent_messages
        assert channel.sent_messages[0].finish_reason == "cancelled"

    asyncio.run(run())


def test_gateway_rejects_channel_manager_and_channels_together() -> None:
    async def run() -> None:
        bus = MessageBus()
        try:
            await run_gateway(
                service=FakeService(),
                manage_service_lifecycle=False,
                bus=bus,
                channel_manager=ChannelManager(bus),
                channels=[MemoryChannelAdapter(bus)],
                stop_event=asyncio.Event(),
            )
        except ValueError as exc:
            assert "either channel_manager or channels" in str(exc)
        else:
            raise AssertionError("expected ValueError")

    asyncio.run(run())


def test_gateway_fails_fast_for_service_without_handle_inbound_message() -> None:
    async def run() -> None:
        try:
            await run_gateway(
                service=InvalidService(),
                manage_service_lifecycle=False,
                bus=MessageBus(),
                stop_event=asyncio.Event(),
            )
        except TypeError as exc:
            assert "handle_inbound_message" in str(exc)
        else:
            raise AssertionError("expected TypeError")

    asyncio.run(run())


def test_agent_service_maps_inbound_error_to_structured_outbound() -> None:
    async def run() -> None:
        service = AgentService()

        async def failing_submit_direct(message: str, **kwargs: Any) -> FakeResult:
            raise RuntimeError("boom")

        service.submit_direct = failing_submit_direct  # type: ignore[method-assign]
        outbound = await service.handle_inbound_message(
            InboundMessage(channel="memory", content="hello", session_id="s1", metadata={"source": "test"})
        )

        assert outbound.finish_reason == "error"
        assert outbound.session_id == "s1"
        assert outbound.metadata["error"] == "boom"
        assert outbound.metadata["inbound_metadata"] == {"source": "test"}

    asyncio.run(run())


def test_agent_service_maps_stopped_runtime_to_stopped_outbound() -> None:
    async def run() -> None:
        service = AgentService()

        async def stopped_submit_direct(message: str, **kwargs: Any) -> FakeResult:
            raise RuntimeError("AgentLoop.submit_direct() is not accepting new tasks after stop()")

        service.submit_direct = stopped_submit_direct  # type: ignore[method-assign]
        outbound = await service.handle_inbound_message(
            InboundMessage(channel="memory", content="hello", session_id="s1")
        )

        assert outbound.finish_reason == "stopped"
        assert "not accepting new tasks" in outbound.metadata["error"]

    asyncio.run(run())


def test_channel_manager_start_cancellation_rolls_back_started_channels() -> None:
    class StartedChannel:
        name = "started"

        def __init__(self, bus: MessageBus) -> None:
            self.bus = bus
            self.stopped = False

        async def start(self) -> None:
            pass

        async def stop(self) -> None:
            self.stopped = True

        async def send(self, message: Any) -> None:
            pass

    class BlockingChannel:
        name = "blocking"

        def __init__(self, bus: MessageBus) -> None:
            self.bus = bus
            self.entered = asyncio.Event()

        async def start(self) -> None:
            self.entered.set()
            await asyncio.sleep(10)

        async def stop(self) -> None:
            pass

        async def send(self, message: Any) -> None:
            pass

    async def run() -> None:
        bus = MessageBus()
        started = StartedChannel(bus)
        blocking = BlockingChannel(bus)
        manager = ChannelManager(bus)
        manager.register(started)
        manager.register(blocking)

        task = asyncio.create_task(manager.start())
        await blocking.entered.wait()
        task.cancel()
        try:
            await task
        except asyncio.CancelledError:
            pass
        else:
            raise AssertionError("expected cancellation")

        assert started.stopped

    asyncio.run(run())