qtalker---/backend/tests/test_ws_session.py

import asyncio

import pytest
from sqlalchemy import select

import app.ws
from app.entities import fallback_profile
from app.models.contact_session import ContactSession
from app.models.event import Event
from app.rate_limit import RateLimiter


class FakeSpiritService:
    async def mint_profile(self, signature, channel, anomalies, language="en"):
        return fallback_profile(signature)

    async def fragment(self, source, anomaly, language="en"):
        return "listen"

    async def wire_whisper(self, telemetry, language="en"):
        return "the wire hums"

    def chat_stream(self, entity, question, history, language="en"):
        async def gen():
            for token in ["I ", "am ", "here."]:
                yield token

        return gen()

    def ambient_ready(self):
        return False


async def _fake_synth(text, voice, profile):
    return b"RIFFfake wav bytes"


@pytest.fixture(autouse=True)
def _fake_spirits(monkeypatch):
    monkeypatch.setattr(app.ws, "spirit_service", FakeSpiritService())
    monkeypatch.setattr(app.ws, "synthesize_spirit_voice", _fake_synth)


def _read_until(ws, msg_type, max_frames=30, **match):
    for _ in range(max_frames):
        frame = ws.receive_json()
        if frame.get("type") != msg_type:
            continue
        if all(frame.get(key) == value for key, value in match.items()):
            return frame
    raise AssertionError(f"never saw frame of type {msg_type!r} matching {match!r}")


def _login(sync_client, username="wsmedium"):
    sync_client.post("/auth/register", json={"username": username, "password": "spookyspooky"})
    sync_client.post("/auth/login", json={"username": username, "password": "spookyspooky"})
    return sync_client.cookies.get("qm_session")


def _ws_connect(sync_client, token):
    # TestClient upgrades over ws:// (insecure), so the jar withholds the
    # Secure qm_session cookie. Pass it explicitly — real browsers on https
    # send it on the upgrade automatically.
    return sync_client.websocket_connect(
        "/ws/session", headers={"cookie": f"qm_session={token}"}
    )


def test_websocket_requires_authentication(sync_client):
    with pytest.raises(Exception):
        with sync_client.websocket_connect("/ws/session"):
            pass


@pytest.mark.asyncio
async def test_websocket_ping_pong_and_session_lifecycle(sync_client, db_session):
    _login(sync_client)

    with _ws_connect(sync_client, sync_client.cookies.get("qm_session")) as ws:
        _read_until(ws, "session")
        ws.send_json({"type": "ping"})
        assert _read_until(ws, "pong") == {"type": "pong"}

        sessions = (await db_session.execute(select(ContactSession))).scalars().all()
        assert len(sessions) == 1
        assert sessions[0].ended_at is None

    # The server marks the session ended in its disconnect handler; give the
    # portal loop a moment to commit before asserting.
    for _ in range(40):
        await asyncio.sleep(0.05)
        db_session.expire_all()
        sessions = (await db_session.execute(select(ContactSession))).scalars().all()
        if sessions[0].ended_at is not None:
            break
    assert sessions[0].ended_at is not None


@pytest.mark.asyncio
async def test_summon_mints_entity_and_greets(sync_client, db_session):
    _login(sync_client, "summoner")

    with _ws_connect(sync_client, sync_client.cookies.get("qm_session")) as ws:
        _read_until(ws, "session")
        ws.send_json({"type": "summon"})
        entity_frame = _read_until(ws, "entity")
        assert entity_frame["is_new"] is True
        assert entity_frame["entity"]["name"]
        assert entity_frame["entity"]["rarity"] in ("common", "uncommon", "rare", "mythic")
        assert entity_frame["entity"]["voice"]["voice_id"]
        greeting = _read_until(ws, "utterance")
        assert greeting["kind"] == "greeting"
        assert greeting["text"]

    sessions = (await db_session.execute(select(ContactSession))).scalars().all()
    assert sessions[0].entity_id is not None


@pytest.mark.asyncio
async def test_question_streams_reply_and_records_history(sync_client, db_session):
    _login(sync_client, "seeker")

    with _ws_connect(sync_client, sync_client.cookies.get("qm_session")) as ws:
        _read_until(ws, "session")
        ws.send_json({"type": "question", "text": "Are you at peace?"})
        _read_until(ws, "entity")  # auto-summoned before answering
        _read_until(ws, "reply_start")
        reply_end = _read_until(ws, "reply_end")
        assert reply_end["text"] == "I am here."

    events = (await db_session.execute(select(Event).order_by(Event.created_at))).scalars().all()
    kinds = [event.kind for event in events]
    assert "question" in kinds
    assert "reply" in kinds
    reply_event = next(event for event in events if event.kind == "reply")
    assert reply_event.text == "I am here."


@pytest.mark.asyncio
async def test_anomalies_attune_then_produce_fragments(sync_client):
    _login(sync_client, "listener")

    with _ws_connect(sync_client, sync_client.cookies.get("qm_session")) as ws:
        _read_until(ws, "session")
        for i in range(3):
            ws.send_json(
                {"type": "anomaly", "source": "radio", "frequency": 101.1 + i, "magnitude": 6.5}
            )
        entity_frame = _read_until(ws, "entity")
        assert entity_frame["is_new"] is True

        ws.send_json({"type": "anomaly", "source": "radio", "frequency": 104.0, "magnitude": 7.1})
        fragment = _read_until(ws, "utterance", kind="fragment")
        assert fragment["text"] == "listen"


@pytest.mark.asyncio
async def test_same_signature_recontacts_same_entity(sync_client):
    _login(sync_client, "mediumx")
    anomalies = [
        {"type": "anomaly", "source": "radio", "frequency": 101.0 + i, "magnitude": 5.0 + i}
        for i in range(4)
    ]

    with _ws_connect(sync_client, sync_client.cookies.get("qm_session")) as ws:
        _read_until(ws, "session")
        for anomaly in anomalies:
            ws.send_json(anomaly)
        first = _read_until(ws, "entity")["entity"]["name"]

    with _ws_connect(sync_client, sync_client.cookies.get("qm_session")) as ws:
        _read_until(ws, "session")
        for anomaly in anomalies:
            ws.send_json(anomaly)
        second_frame = _read_until(ws, "entity")
        assert second_frame["entity"]["name"] == first
        assert second_frame["is_new"] is False
        assert second_frame["entity"]["contact_count"] == 2


@pytest.mark.asyncio
async def test_summon_rate_limited_per_account(sync_client, monkeypatch):
    # Swap in a tight, test-scoped limiter so this doesn't depend on (or
    # pollute) the shared module-level budget other tests draw from.
    monkeypatch.setattr(
        app.ws, "summon_limiter", RateLimiter(max_requests=2, window_seconds=60)
    )
    _login(sync_client, "account-limited")

    with _ws_connect(sync_client, sync_client.cookies.get("qm_session")) as ws:
        _read_until(ws, "session")
        for _ in range(2):
            ws.send_json({"type": "summon"})
            _read_until(ws, "entity")
            _read_until(ws, "utterance", kind="greeting")

        ws.send_json({"type": "summon"})
        rejection = _read_until(ws, "error")
        assert rejection["code"] == "rate_limited"
        assert rejection["message"] == (
            "The veil is crowded. The spirits need a moment before another summoning."
        )


@pytest.mark.asyncio
async def test_summon_rate_limited_per_ip_even_with_fresh_account(
    sync_client, monkeypatch
):
    # Starve only the IP bucket; the per-account limiter stays at its
    # production default so each account below has plenty of its own budget
    # left. This proves the IP limiter alone can reject a request — the
    # gap the per-account-only limiters left open.
    monkeypatch.setattr(
        app.ws, "summon_ip_limiter", RateLimiter(max_requests=1, window_seconds=60)
    )

    _login(sync_client, "ip-limited-a")
    with _ws_connect(sync_client, sync_client.cookies.get("qm_session")) as ws:
        _read_until(ws, "session")
        ws.send_json({"type": "summon"})
        _read_until(ws, "entity")  # spends the single per-IP slot

    # A different account — its own per-account budget is untouched — but
    # every connection in this test shares the same (simulated) source IP,
    # which is already spent.
    _login(sync_client, "ip-limited-b")
    with _ws_connect(sync_client, sync_client.cookies.get("qm_session")) as ws:
        _read_until(ws, "session")
        ws.send_json({"type": "summon"})
        rejection = _read_until(ws, "error")
        assert rejection["code"] == "rate_limited"
        assert rejection["message"] == (
            "The veil is crowded. The spirits need a moment before another summoning."
        )


@pytest.mark.asyncio
async def test_summon_per_ip_bucket_follows_cf_connecting_ip_not_socket_peer(
    sync_client, monkeypatch
):
    # Every connection in this test suite shares the same simulated socket
    # peer (TestClient has no real network). Without preferring
    # CF-Connecting-IP, distinct visitors behind the Cloudflare Tunnel would
    # collapse into one shared per-IP bucket — this proves they don't.
    monkeypatch.setattr(
        app.ws, "summon_ip_limiter", RateLimiter(max_requests=1, window_seconds=60)
    )

    token_a = _login(sync_client, "cf-ip-a")
    with sync_client.websocket_connect(
        "/ws/session",
        headers={"cookie": f"qm_session={token_a}", "cf-connecting-ip": "203.0.113.1"},
    ) as ws:
        _read_until(ws, "session")
        ws.send_json({"type": "summon"})
        _read_until(ws, "entity")  # spends visitor A's per-IP slot only

    token_b = _login(sync_client, "cf-ip-b")
    with sync_client.websocket_connect(
        "/ws/session",
        headers={"cookie": f"qm_session={token_b}", "cf-connecting-ip": "203.0.113.2"},
    ) as ws:
        _read_until(ws, "session")
        ws.send_json({"type": "summon"})
        # A different visitor IP behind the same tunnel — must not be
        # rejected by visitor A's already-spent bucket.
        _read_until(ws, "entity")