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

import asyncio
import uuid

import pytest
from sqlalchemy import select

import app.ws
from app.entities import fallback_profile
from app.inventory import SUMMON_ESSENCE_TRICKLE
from app.models.contact_session import ContactSession
from app.models.event import Event
from app.models.inventory_item import InventoryItem
from app.models.user import User
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, instability=0.0):
    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
        reply_start = _read_until(ws, "reply_start")
        assert 0.05 <= reply_start["stability"] <= 0.98
        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")


@pytest.mark.asyncio
async def test_summon_credits_essence_trickle(sync_client, db_session, monkeypatch):
    # The module-level summon_ip_limiter/summon_limiter are shared singletons
    # that persist real hit counts across every test in this file (they're
    # only overridden where a test explicitly monkeypatches them, as earlier
    # tests above do) — fresh, generous instances here keep this test from
    # depending on how many summons ran before it.
    monkeypatch.setattr(app.ws, "summon_ip_limiter", RateLimiter(max_requests=1000, window_seconds=60))
    monkeypatch.setattr(app.ws, "summon_limiter", RateLimiter(max_requests=1000, window_seconds=60))

    token = _login(sync_client, "trickle-earner")
    user_id = uuid.UUID(
        sync_client.get("/auth/me", headers={"cookie": f"qm_session={token}"}).json()["id"]
    )

    with _ws_connect(sync_client, token) as ws:
        _read_until(ws, "session")
        ws.send_json({"type": "summon"})
        _read_until(ws, "entity")
        _read_until(ws, "utterance", kind="greeting")

    for _ in range(40):
        db_session.expire_all()
        user = await db_session.get(User, user_id)
        if user.essence == SUMMON_ESSENCE_TRICKLE:
            break
        await asyncio.sleep(0.05)
    assert user.essence == SUMMON_ESSENCE_TRICKLE


@pytest.mark.asyncio
async def test_summon_high_rarity_item_drop_is_persisted_and_sent(
    sync_client, db_session, monkeypatch
):
    # Force a hit so the wiring is exercised deterministically rather than
    # relying on the real roll odds (those are covered statistically in
    # test_inventory.py).
    forced_item = {
        "item_type": "curio",
        "item_key": "veil_thread",
        "payload": {"trigger": "summon_mythic"},
    }
    monkeypatch.setattr(app.ws, "roll_item_drop", lambda trigger: forced_item)
    # See test_summon_credits_essence_trickle above re: why this is reset.
    monkeypatch.setattr(app.ws, "summon_ip_limiter", RateLimiter(max_requests=1000, window_seconds=60))
    monkeypatch.setattr(app.ws, "summon_limiter", RateLimiter(max_requests=1000, window_seconds=60))

    class MythicSpiritService:
        async def mint_profile(self, signature, channel, anomalies, language="en"):
            profile = fallback_profile(signature)
            profile["rarity"] = "mythic"
            return profile

        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():
                yield "hi"

            return gen()

        def ambient_ready(self):
            return False

    monkeypatch.setattr(app.ws, "spirit_service", MythicSpiritService())

    token = _login(sync_client, "jackpot-seeker")
    user_id = uuid.UUID(
        sync_client.get("/auth/me", headers={"cookie": f"qm_session={token}"}).json()["id"]
    )

    with _ws_connect(sync_client, token) as ws:
        _read_until(ws, "session")
        ws.send_json({"type": "summon"})
        entity_frame = _read_until(ws, "entity")
        assert entity_frame["entity"]["rarity"] == "mythic"
        drop_frame = _read_until(ws, "item_drop")
        assert drop_frame["item"] == forced_item

    for _ in range(40):
        db_session.expire_all()
        items = (
            await db_session.execute(
                select(InventoryItem).where(InventoryItem.user_id == user_id)
            )
        ).scalars().all()
        if items:
            break
        await asyncio.sleep(0.05)
    assert len(items) == 1
    assert items[0].item_key == "veil_thread"
    assert items[0].item_type == "curio"


@pytest.mark.asyncio
async def test_summon_common_rarity_does_not_roll_a_drop(sync_client, db_session, monkeypatch):
    # See test_summon_credits_essence_trickle above re: why this is reset.
    monkeypatch.setattr(app.ws, "summon_ip_limiter", RateLimiter(max_requests=1000, window_seconds=60))
    monkeypatch.setattr(app.ws, "summon_limiter", RateLimiter(max_requests=1000, window_seconds=60))

    # Common is the overwhelmingly likely fallback rarity (55% weight) but
    # not guaranteed — force it explicitly so this test can't flake.
    def _common_only(signature: str) -> dict:
        profile = fallback_profile(signature)
        profile["rarity"] = "common"
        return profile

    class CommonSpiritService:
        async def mint_profile(self, signature, channel, anomalies, language="en"):
            return _common_only(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():
                yield "hi"

            return gen()

        def ambient_ready(self):
            return False

    monkeypatch.setattr(app.ws, "spirit_service", CommonSpiritService())
    # If a drop were (incorrectly) rolled for a common-rarity summon, this
    # would make it always hit — proving the rarity gate, not just the odds,
    # is what's preventing a drop here.
    monkeypatch.setattr(
        app.ws, "roll_item_drop", lambda trigger: {"item_type": "x", "item_key": "y", "payload": {}}
    )

    token = _login(sync_client, "commoner")
    user_id = uuid.UUID(
        sync_client.get("/auth/me", headers={"cookie": f"qm_session={token}"}).json()["id"]
    )

    with _ws_connect(sync_client, token) as ws:
        _read_until(ws, "session")
        ws.send_json({"type": "summon"})
        entity_frame = _read_until(ws, "entity")
        assert entity_frame["entity"]["rarity"] == "common"
        greeting = _read_until(ws, "utterance", kind="greeting")
        assert greeting["kind"] == "greeting"

    # Give the (already-awaited, so this should be immediate) reward path a
    # moment to land, then assert no item was ever persisted for this user —
    # the rarity gate, not just the roll odds, is what prevents a drop here.
    for _ in range(20):
        db_session.expire_all()
        items = (
            await db_session.execute(
                select(InventoryItem).where(InventoryItem.user_id == user_id)
            )
        ).scalars().all()
        await asyncio.sleep(0.02)
    assert items == []