PacketStation/tests/server/test_presence.py

"""Phase 8: presence. Who is on which screen and packet, with the role shown; locked sessions drop out and come
back on unlock; revoked sessions are gone at once. Driven through the EventHub directly (as test_events.py does)
and through the API (`/api/events`, `/api/auth/lock`, `/api/auth/unlock`, the idle lock in `resolve`).
"""

from __future__ import annotations

import asyncio
import json
import uuid
from collections.abc import AsyncIterator, Iterator
from typing import Any

import pytest
from fakedir import FAKE_PASSWORD, FakeDirectory
from fastapi.testclient import TestClient
from pshelpers import ORIGIN_HEADERS, needs_pg, sign_in

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

USER_KEYS = {"user_id", "display_name", "role", "screen", "packet_id", "since"}


def _sub(
    hub: events.EventHub,
    name: str,
    *,
    screen: str | None = None,
    packet_id: int | None = None,
    role: str = "Supervisor",
    user_id: Any = None,
    session_id: Any = None,
) -> events.Subscriber:
    return hub.subscribe(
        session_id=session_id or uuid.uuid4(),
        user_id=user_id or uuid.uuid4(),
        display_name=name,
        screen=screen,
        packet_id=packet_id,
        role=role,
    )


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 _on(hub: events.EventHub, packet_id: int) -> list[tuple[str, str, str | None]]:
    return [(u["display_name"], u["role"], u["screen"]) for u in hub.presence()["users"] if u["packet_id"] == packet_id]


def _names(hub: events.EventHub) -> list[str]:
    return [u["display_name"] for u in hub.presence()["users"]]


def _session(hub: events.EventHub, sub: events.Subscriber, action: str) -> None:
    hub._deliver(json.dumps({"event": "session", "session_id": sub.session_id, "action": action}))


# ================================================================ the hub


def test_two_users_on_one_packet_see_each_other() -> None:
    async def run() -> None:
        hub = events.EventHub()
        wes = _sub(hub, "Wes Super", screen="supervisor", packet_id=7, role="Supervisor")
        andy = _sub(hub, "Andy Welder", screen="stall", packet_id=7, role="Weld Stall")
        _sub(hub, "Clara Clerk", screen="clerical", role="Clerical")
        users = hub.presence()["users"]
        assert all(set(u) == USER_KEYS for u in users), users
        assert all(u["since"] for u in users)
        assert sorted(_on(hub, 7)) == [
            ("Andy Welder", "Weld Stall", "stall"),
            ("Wes Super", "Supervisor", "supervisor"),
        ]
        last = (await _drain(wes.queue))[-1]
        assert last[0] == "presence" and {u["display_name"] for u in last[1]["users"]} == {
            "Wes Super",
            "Andy Welder",
            "Clara Clerk",
        }
        assert (await _drain(andy.queue))[-1] == last

    asyncio.run(run())


def test_one_entry_per_user_screen_and_packet() -> None:
    async def run() -> None:
        hub = events.EventHub()
        uid = uuid.uuid4()
        _sub(hub, "Wes Super", screen="supervisor", packet_id=7, user_id=uid)
        _sub(hub, "Wes Super", screen="supervisor", packet_id=7, user_id=uid)  # a second tab
        _sub(hub, "Wes Super", screen="supervisor", packet_id=8, user_id=uid)
        assert len(_on(hub, 7)) == 1 and len(_on(hub, 8)) == 1
        assert _names(hub) == ["Wes Super", "Wes Super"]

    asyncio.run(run())


def test_leaving_drops_out() -> None:
    async def run() -> None:
        hub = events.EventHub()
        wes = _sub(hub, "Wes Super", packet_id=7)
        andy = _sub(hub, "Andy Welder", packet_id=7, role="Weld Stall")
        await _drain(wes.queue)
        hub.unsubscribe(andy)
        got = await _drain(wes.queue)
        assert got and got[-1][0] == "presence"
        assert [u["display_name"] for u in got[-1][1]["users"]] == ["Wes Super"]
        assert _on(hub, 7) == [("Wes Super", "Supervisor", None)]

    asyncio.run(run())


def test_locked_hides_and_unlocked_brings_back() -> None:
    async def run() -> None:
        hub = events.EventHub()
        wes = _sub(hub, "Wes Super", packet_id=7)
        andy = _sub(hub, "Andy Welder", packet_id=7, role="Weld Stall")
        await _drain(wes.queue)
        await _drain(andy.queue)

        _session(hub, andy, "locked")
        assert _names(hub) == ["Wes Super"]
        to_wes = await _drain(wes.queue)
        assert [n for n, _ in to_wes] == ["presence"]  # the others get presence, never the session event
        assert [u["display_name"] for u in to_wes[-1][1]["users"]] == ["Wes Super"]
        to_andy = await _drain(andy.queue)
        assert ("session", {"session_id": andy.session_id, "action": "locked"}) in to_andy

        _session(hub, andy, "unlocked")
        assert sorted(_names(hub)) == ["Andy Welder", "Wes Super"]
        to_wes = await _drain(wes.queue)
        assert to_wes and to_wes[-1][0] == "presence" and len(to_wes[-1][1]["users"]) == 2
        assert ("session", {"session_id": andy.session_id, "action": "unlocked"}) in await _drain(andy.queue)

    asyncio.run(run())


def test_a_locked_session_hides_only_itself() -> None:
    async def run() -> None:
        hub = events.EventHub()
        uid = uuid.uuid4()
        tablet = _sub(hub, "Wes Super", packet_id=7, user_id=uid)
        _sub(hub, "Wes Super", packet_id=7, user_id=uid)  # the same person on a second device
        _session(hub, tablet, "locked")
        assert _on(hub, 7) == [("Wes Super", "Supervisor", None)]

    asyncio.run(run())


def test_revoked_is_gone_at_once() -> None:
    async def run() -> None:
        hub = events.EventHub()
        wes = _sub(hub, "Wes Super", packet_id=7)
        andy = _sub(hub, "Andy Welder", packet_id=7, role="Weld Stall")
        await _drain(wes.queue)
        _session(hub, andy, "revoked")
        assert _names(hub) == ["Wes Super"]  # before the stream has ended
        got = await _drain(wes.queue)
        assert got and got[-1][0] == "presence" and len(got[-1][1]["users"]) == 1
        assert ("session", {"session_id": andy.session_id, "action": "revoked"}) in await _drain(andy.queue)

    asyncio.run(run())


def test_stream_forwards_lock_and_keeps_going() -> None:
    async def run() -> list[str]:
        hub = events.EventHub()
        sub = _sub(hub, "Wes Super", screen="supervisor")
        await _drain(sub.queue)

        async def connected() -> bool:
            return False

        gen = events.stream(hub, sub, hello={"session_id": sub.session_id}, is_disconnected=connected, ping_seconds=5)
        chunks = [await gen.__anext__()]
        _session(hub, sub, "locked")
        for _ in range(3):
            chunks.append(await asyncio.wait_for(gen.__anext__(), 2))
            if chunks[-1].startswith("event: session"):
                break
        _session(hub, sub, "revoked")
        rest = [c async for c in gen]
        assert hub.presence() == {"users": []}
        return chunks + rest

    chunks = asyncio.run(run())
    assert chunks[0].startswith("event: hello\n")
    names = [c.split("\n", 1)[0] for c in chunks[1:]]
    assert "event: session" in names  # the lock reached the locked session's own stream
    assert any('"action":"locked"' in c for c in chunks)


def test_presence_by_screen() -> None:
    async def run() -> None:
        hub = events.EventHub()
        _sub(hub, "Wes Super", screen="supervisor")
        _sub(hub, "Lou Lead", screen="supervisor")
        _sub(hub, "Quinn QC", screen="qc", role="Q.C.")
        locked = _sub(hub, "Andy Welder", screen="stall", role="Weld Stall")
        _session(hub, locked, "locked")
        by = hub.presence_by_screen()
        assert sorted(u["display_name"] for u in by["supervisor"]) == ["Lou Lead", "Wes Super"]
        assert [u["display_name"] for u in by["qc"]] == ["Quinn QC"]
        assert not by.get("stall")

    asyncio.run(run())


# ================================================================ through the API


@pytest.fixture
def client(pg_settings: Settings, fake_dir: FakeDirectory, dev_visual: None) -> Iterator[TestClient]:
    from packetstation.main import create_app

    with TestClient(create_app(pg_settings, directory=fake_dir), headers=ORIGIN_HEADERS) as c:
        yield c


def _signed(client: TestClient, sam: str) -> tuple[TestClient, str]:
    c = TestClient(client.app, headers=ORIGIN_HEADERS)
    r = sign_in(c, sam)
    assert r.status_code == 200, r.text
    return c, str(r.json()["session"]["id"])


@needs_pg
def test_events_stream_carries_the_role(client: TestClient, monkeypatch: pytest.MonkeyPatch) -> None:
    """The stream body is replaced by one that answers the hub's presence and ends, so the request returns."""
    from packetstation.api import events as events_api

    async def snapshot(
        hub: events.EventHub, sub: events.Subscriber, *, hello: dict[str, Any], **_: Any
    ) -> AsyncIterator[str]:
        try:
            yield events.format_sse("presence", hub.presence())
        finally:
            hub.unsubscribe(sub)

    monkeypatch.setattr(events_api, "stream", snapshot)

    def users(c: TestClient, **params: Any) -> list[dict[str, Any]]:
        r = c.get("/api/events", params=params)
        assert r.status_code == 200, r.text
        line = next(x for x in r.text.splitlines() if x.startswith("data: "))
        out: list[dict[str, Any]] = json.loads(line[6:])["users"]
        return out

    wes, _ = _signed(client, "sup.wes")
    me = users(wes, screen="supervisor", packet_id=5)
    assert [(u["display_name"], u["role"], u["screen"], u["packet_id"]) for u in me] == [
        ("Wes Super", "Supervisor", "supervisor", 5)
    ]
    assert me[0]["since"]
    # a screen the user does not have: no screen, the first role's label
    other = users(wes, screen="qc")
    assert [(u["role"], u["screen"]) for u in other] == [("Supervisor", None)]
    for sam, screen, role in (
        ("pm.paula", "pm", "Project Mgmt"),
        ("weld.andy", "stall", "Weld Stall"),
        ("clerk.clara", "clerical", "Clerical"),
        ("wh.wade", "warehouse", "Warehouse"),
        ("qc.quinn", "qc", "Q.C."),
        ("admin.ann", "admin", "Admin"),
    ):
        c, _ = _signed(client, sam)
        assert [u["role"] for u in users(c, screen=screen)] == [role], sam
    paula, _ = _signed(client, "pm.paula")
    assert [u["role"] for u in users(paula)] == ["Project Mgmt"]  # no screen: the primary role


async def _wait_session(q: asyncio.Queue[events.Item], action: str) -> dict[str, Any]:
    while True:
        name, data = await asyncio.wait_for(q.get(), timeout=10)
        if name == "session" and data.get("action") == action:
            return data


@needs_pg
def test_lock_and_unlock_publish_session_events(client: TestClient, pg_settings: Settings) -> None:
    wes, sid = _signed(client, "sup.wes")

    async def run() -> None:
        hub = events.EventHub(connect=lambda: db.connect_as("app", pg_settings, autocommit=True))
        sub = _sub(hub, "Wes Super", screen="supervisor", packet_id=3, session_id=sid)
        other = _sub(hub, "Andy Welder", screen="stall", packet_id=3, role="Weld Stall")
        hub.start()
        try:
            for _ in range(100):
                if hub.connected:
                    break
                await asyncio.sleep(0.05)
            assert hub.connected
            r = await asyncio.to_thread(wes.post, "/api/auth/lock")
            assert r.status_code == 204, r.text
            data = await _wait_session(sub.queue, "locked")
            assert data["session_id"] == sid
            assert _names(hub) == ["Andy Welder"]
            assert other.queue.qsize() >= 1  # the others got a presence update
            r = await asyncio.to_thread(
                wes.post, "/api/auth/unlock", json={"factor": "password", "secret": FAKE_PASSWORD}
            )
            assert r.status_code == 200, r.text
            await _wait_session(sub.queue, "unlocked")
            assert sorted(_names(hub)) == ["Andy Welder", "Wes Super"]
        finally:
            await asyncio.to_thread(hub.stop)

    asyncio.run(run())


@needs_pg
def test_idle_lock_publishes(client: TestClient, pg_settings: Settings) -> None:
    wes, sid = _signed(client, "sup.wes")

    async def run() -> None:
        hub = events.EventHub(connect=lambda: db.connect_as("app", pg_settings, autocommit=True))
        sub = _sub(hub, "Wes Super", session_id=sid)
        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:
                conn.execute("UPDATE sessions SET last_seen_at = now() - interval '5 minutes' WHERE id = %s", (sid,))
            r = await asyncio.to_thread(wes.get, "/api/auth/me")
            assert r.status_code == 200 and r.json()["session"]["locked"] is True, r.text
            await _wait_session(sub.queue, "locked")
            assert _names(hub) == []
        finally:
            await asyncio.to_thread(hub.stop)

    asyncio.run(run())