PacketStation/tests/server/test_events.py

from __future__ import annotations

import asyncio
import json
import uuid

from pshelpers import needs_pg

from packetstation import db, events
from packetstation.settings import Settings


def _sub(hub: events.EventHub, name: str, screen: str | None = None) -> events.Subscriber:
    return hub.subscribe(session_id=uuid.uuid4(), user_id=uuid.uuid4(), display_name=name, screen=screen)


async def _drain(q: asyncio.Queue[events.Item]) -> list[events.Item]:
    await asyncio.sleep(0)  # let call_soon_threadsafe callbacks run
    out = []
    while not q.empty():
        out.append(q.get_nowait())
    return out


def test_fan_out_presence_and_session_targeting() -> None:
    async def run() -> None:
        hub = events.EventHub()
        a = _sub(hub, "Ann", "admin")
        b = _sub(hub, "Bo", "stall")
        got_a = await _drain(a.queue)
        assert [n for n, _ in got_a] == ["presence", "presence"]
        assert got_a[-1][1] == {
            "users": [
                {"user_id": a.user_id, "display_name": "Ann", "screen": "admin", "packet_id": None},
                {"user_id": b.user_id, "display_name": "Bo", "screen": "stall", "packet_id": None},
            ]
        }
        await _drain(b.queue)

        hub._deliver(json.dumps({"event": "changed", "collection": "packets", "id": 7}))
        hub._deliver(json.dumps({"collection": "projects"}))  # no event name means changed
        hub._deliver("not json")
        for q in (a.queue, b.queue):
            assert await _drain(q) == [
                ("changed", {"collection": "packets", "id": 7}),
                ("changed", {"collection": "projects"}),
            ]

        hub._deliver(json.dumps({"event": "session", "session_id": b.session_id, "action": "revoked"}))
        assert await _drain(a.queue) == []
        assert await _drain(b.queue) == [("session", {"session_id": b.session_id, "action": "revoked"})]

        hub.unsubscribe(b)
        got = await _drain(a.queue)
        users = [{"user_id": a.user_id, "display_name": "Ann", "screen": "admin", "packet_id": None}]
        assert got == [("presence", {"users": users})]

    asyncio.run(run())


def test_sse_stream_hello_ping_and_end_on_own_revoke() -> None:
    async def run() -> list[str]:
        hub = events.EventHub()
        sub = _sub(hub, "Ann")
        await _drain(sub.queue)

        async def connected() -> bool:
            return False

        gen = events.stream(
            hub, sub, hello={"session_id": sub.session_id, "build": "t"}, is_disconnected=connected, ping_seconds=0.05
        )
        chunks = [await gen.__anext__(), await gen.__anext__()]
        hub._deliver(json.dumps({"event": "changed", "collection": "packets"}))
        chunks.append(await gen.__anext__())
        hub._deliver(json.dumps({"event": "session", "session_id": sub.session_id, "action": "revoked"}))
        chunks.append(await gen.__anext__())
        rest = [c async for c in gen]
        assert rest == []
        assert hub.presence() == {"users": []}
        return chunks

    chunks = asyncio.run(run())
    assert chunks[0].startswith("event: hello\ndata: {") and '"build":"t"' in chunks[0]
    assert chunks[1] == ": ping\n\n"
    assert chunks[2] == 'event: changed\ndata: {"collection":"packets"}\n\n'
    assert chunks[3].startswith("event: session\n")


@needs_pg
def test_pg_notify_round_trip(pg_settings: Settings) -> None:
    async def run() -> events.Item:
        hub = events.EventHub(connect=lambda: db.connect_as("app", pg_settings, autocommit=True))
        sub = _sub(hub, "Ann")
        await _drain(sub.queue)
        hub.start()
        try:
            for _ in range(100):
                if hub.connected:
                    break
                await asyncio.sleep(0.05)
            assert hub.connected
            with db.connect_as("app", pg_settings, autocommit=True) as conn:
                events.publish(conn, {"event": "changed", "collection": "packets", "id": "7A"})
            return await asyncio.wait_for(sub.queue.get(), timeout=5)
        finally:
            await asyncio.to_thread(hub.stop)

    assert asyncio.run(run()) == ("changed", {"collection": "packets", "id": "7A"})