plm-lottery/tests/unit/test_round_events.py

import asyncio

import pytest

from app.rounds.events import EVICTED, RoundEventBroadcaster, RoundEventCapacityError


async def test_publish_wakes_up_subscriber():
    broadcaster = RoundEventBroadcaster()
    queue = broadcaster.subscribe()

    broadcaster.publish()

    await asyncio.wait_for(queue.get(), timeout=1)


async def test_publish_with_no_subscribers_does_not_raise():
    broadcaster = RoundEventBroadcaster()
    broadcaster.publish()  # no subscribers yet — must be a no-op, not an error


async def test_publish_coalesces_when_subscriber_has_not_drained():
    """The queue is maxsize=1: a second publish() before the subscriber reads
    the first notification must not block or raise — it's fine to drop it,
    since the subscriber will refetch full state anyway on the first one."""
    broadcaster = RoundEventBroadcaster()
    queue = broadcaster.subscribe()

    broadcaster.publish()
    broadcaster.publish()  # would raise QueueFull if not guarded

    assert queue.qsize() == 1


async def test_unsubscribe_stops_delivery():
    broadcaster = RoundEventBroadcaster()
    queue = broadcaster.subscribe()
    broadcaster.unsubscribe(queue)

    broadcaster.publish()

    assert queue.empty()


async def test_subscribe_rejects_past_the_cap():
    broadcaster = RoundEventBroadcaster(max_subscribers=3)
    for _ in range(3):
        broadcaster.subscribe()

    with pytest.raises(RoundEventCapacityError):
        broadcaster.subscribe()


async def test_unsubscribe_frees_a_capacity_slot():
    broadcaster = RoundEventBroadcaster(max_subscribers=1)
    queue = broadcaster.subscribe()

    with pytest.raises(RoundEventCapacityError):
        broadcaster.subscribe()

    broadcaster.unsubscribe(queue)
    broadcaster.subscribe()  # no longer at capacity


# --- B-38: a single IP must not be able to exhaust the global cap and degrade
# every other user to polling. ----------------------------------------------------


async def test_subscribe_evicts_the_same_ips_oldest_connection_past_its_cap():
    """Past MAX_SUBSCRIBERS_PER_IP, one more stream from the *same* IP evicts
    that IP's own oldest connection rather than being refused — bounds one
    source's footprint without an outright block."""
    broadcaster = RoundEventBroadcaster(max_subscribers=100, max_per_ip=2)
    first = broadcaster.subscribe("1.2.3.4")
    second = broadcaster.subscribe("1.2.3.4")

    third = broadcaster.subscribe("1.2.3.4")  # past the per-IP cap of 2

    assert await asyncio.wait_for(first.get(), timeout=1) is EVICTED
    assert second.empty()  # untouched — only the oldest was evicted
    assert third is not None


async def test_subscribe_does_not_evict_across_different_ips():
    """A different IP hitting its own cap must never evict an unrelated IP's
    connection — that would let one abusive source crowd out real users, which
    is exactly what the global-only cap used to allow."""
    broadcaster = RoundEventBroadcaster(max_subscribers=100, max_per_ip=1)
    other_ip_queue = broadcaster.subscribe("9.9.9.9")

    broadcaster.subscribe("1.2.3.4")
    broadcaster.subscribe("1.2.3.4")  # evicts 1.2.3.4's own oldest, not 9.9.9.9's

    assert other_ip_queue.empty()


async def test_subscribe_still_enforces_the_global_cap_across_many_ips():
    """The per-IP cap doesn't replace the global backstop — spreading across
    enough distinct IPs must still eventually hit MAX_SUBSCRIBERS."""
    broadcaster = RoundEventBroadcaster(max_subscribers=3, max_per_ip=1)
    broadcaster.subscribe("1.1.1.1")
    broadcaster.subscribe("2.2.2.2")
    broadcaster.subscribe("3.3.3.3")

    with pytest.raises(RoundEventCapacityError):
        broadcaster.subscribe("4.4.4.4")


async def test_unsubscribe_clears_the_per_ip_tracking_too():
    """Regression guard: unsubscribe must forget the queue's IP association, or
    a churned-through connection would keep counting against that IP's cap
    forever."""
    broadcaster = RoundEventBroadcaster(max_subscribers=100, max_per_ip=1)
    queue = broadcaster.subscribe("1.2.3.4")
    broadcaster.unsubscribe(queue)

    broadcaster.subscribe("1.2.3.4")  # must not evict anything — nothing left to evict
    assert queue.empty()


async def test_subscribe_defaults_to_a_shared_ip_when_none_given():
    """Existing callers (and most tests) that don't care about IP isolation
    still share one implicit bucket rather than needing every call updated."""
    broadcaster = RoundEventBroadcaster(max_subscribers=100, max_per_ip=2)
    first = broadcaster.subscribe()
    broadcaster.subscribe()

    broadcaster.subscribe()  # past the default bucket's cap of 2 — evicts, doesn't raise

    assert await asyncio.wait_for(first.get(), timeout=1) is EVICTED