signal-bot/tests/test_dispatcher_concurrency.py

# Copyright (c) 2026 Krzysztof Rudnicki. MIT licence.
"""Frames handled side by side: what still holds, and what each one costs."""

from __future__ import annotations

import asyncio
import datetime as dt
import time
from typing import TYPE_CHECKING

from signal_bot.decorators import cooldown
from signal_bot.dispatcher import Dispatcher
from signal_bot.registry import CommandRegistry
from signal_bot.registry_holder import RegistryHolder
from tests.conftest import envelope

if TYPE_CHECKING:
    from signal_bot.config import Config
    from signal_bot.context import CommandContext
    from signal_bot.discovery import GroupDirectory
    from signal_bot.store import Store
    from signal_bot.transport_fake import FakeTransport

NOW = dt.datetime(2026, 8, 29, 12, 0, tzinfo=dt.UTC)
BOB = ("uuid-bob", "Bob")


async def _settled() -> None:
    for _ in range(3):
        await asyncio.sleep(0)


def _dispatcher(
    config: Config,
    directory: GroupDirectory,
    store: Store,
    transport: FakeTransport,
    gate: asyncio.Event,
    runs: list[str],
) -> Dispatcher:
    """A ``!cat`` that blocks on ``gate`` before answering, with a cooldown."""

    async def handler(ctx: CommandContext) -> None:
        runs.append(ctx.message.source_uuid)
        await gate.wait()
        await ctx.reply("meow")

    registry = CommandRegistry()
    registry.register("cat", ["!cat"], cooldown(10, "wait {seconds}")(handler), "kot")
    return Dispatcher(
        config=config,
        directory=directory,
        holder=RegistryHolder(lambda: registry),
        store=store,
        transport=transport,
        clock=lambda: NOW,
    )


async def test_the_same_persons_second_command_waits_and_is_then_refused(
    config: Config, directory: GroupDirectory, store: Store, transport: FakeTransport
) -> None:
    """Two ``!cat`` in one second used to both pass the gate; now one does.

    Without the per-person lock, both would read "never used" before either
    had recorded a use, and the cooldown would be a decoration.
    """
    gate = asyncio.Event()
    runs: list[str] = []
    dispatcher = _dispatcher(config, directory, store, transport, gate, runs)

    first = asyncio.ensure_future(dispatcher.handle(envelope(text="!cat")))
    second = asyncio.ensure_future(dispatcher.handle(envelope(text="!cat")))
    await _settled()
    assert runs == ["uuid-alice"]
    assert not second.done()

    gate.set()
    served_first, served_second = await asyncio.gather(first, second)
    assert runs == ["uuid-alice"]
    assert served_first is not None
    assert served_second is not None
    assert (served_first.outcome, served_second.outcome) == ("ok", "refused_cooldown")
    assert transport.texts[0] == "meow"
    assert transport.texts[1].startswith("wait ")


async def test_another_persons_command_is_answered_while_the_first_still_runs(
    config: Config, directory: GroupDirectory, store: Store, transport: FakeTransport
) -> None:
    gate = asyncio.Event()
    runs: list[str] = []
    dispatcher = _dispatcher(config, directory, store, transport, gate, runs)

    alice = asyncio.ensure_future(dispatcher.handle(envelope(text="!cat")))
    bob = asyncio.ensure_future(dispatcher.handle(envelope(text="!cat", sender=BOB)))
    await _settled()
    # Both are inside the handler at once: Bob did not queue behind Alice.
    assert runs == ["uuid-alice", "uuid-bob"]
    gate.set()
    await asyncio.gather(alice, bob)
    assert len(transport.sent) == 2
    assert dispatcher.users.held == frozenset()


async def test_what_a_served_command_reports(
    config: Config, directory: GroupDirectory, store: Store, transport: FakeTransport
) -> None:
    gate = asyncio.Event()
    gate.set()
    dispatcher = _dispatcher(config, directory, store, transport, gate, [])

    served = await dispatcher.handle(envelope(text="!cat"), received_at=time.perf_counter() - 0.25)

    assert served is not None
    assert (served.command, served.outcome) == ("cat", "ok")
    # A quarter of a second sat in the queue before this frame got a slot.
    assert 200 <= served.queue_ms <= 400
    assert served.dispatch_ms >= 0
    assert served.handler_ms >= served.send_ms
    assert served.upstream_ms == 0


async def test_a_frame_with_no_command_in_it_costs_nothing_to_report(
    config: Config, directory: GroupDirectory, store: Store, transport: FakeTransport
) -> None:
    dispatcher = _dispatcher(config, directory, store, transport, asyncio.Event(), [])
    assert await dispatcher.handle(envelope(text="hej"), received_at=time.perf_counter()) is None
    assert await dispatcher.handle(envelope(text="!nic")) is None
    assert await dispatcher.handle(envelope(text="!cat", group_id="elsewhere")) is None
    assert dispatcher.users.held == frozenset()


async def test_frames_that_never_reach_a_command_take_no_lock(
    config: Config, directory: GroupDirectory, store: Store, transport: FakeTransport
) -> None:
    """A plain message from Alice must not queue behind Alice's slow ``!cat``."""
    gate = asyncio.Event()
    dispatcher = _dispatcher(config, directory, store, transport, gate, [])
    slow = asyncio.ensure_future(dispatcher.handle(envelope(text="!cat")))
    await _settled()
    assert dispatcher.users.held == {"uuid-alice"}
    plain = await asyncio.wait_for(dispatcher.handle(envelope(text="hello")), timeout=1)
    assert plain is None
    gate.set()
    await slow