logan/server-26
1
0
Fork 0
Code Issues 84 Pull Requests Actions Packages Projects Releases Wiki Activity
Files
e3c68738e8520508019f61f3e4e55c1e430341b9
server-26/drb-c2-core/tests/test_discord_leases.py
T
Logan CusanoandClaude Opus 5.5 03ba7eacd9 Discord leases: free releasing tokens on gone nodes; stamp before building desired list 
- reconcile_once only walked desired=on streams, so a token left "releasing"
  by /leave while its node was offline (or a deleted node) never returned to
  the pool until that node came back (server-26#204 leak shape). Sweep leased
  tokens whose node is offline/deleted past the lease grace.
- publish_desired stamped release_seq AFTER building the list, so a stop
  landing between the two got a seq whose list still held the stream; an ack
  for it could free a token the node was still told to use. Stamp first.

Three regression tests; each fails on the previous commit.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-04 18:16:31 -04:00

478 lines
20 KiB
Python
Raw Blame History

"""
Radio-bot token leases + Discord desired-state streams
(server-26#203, server-26#204, node-26#25).
The race tests run real concurrent transactions against tests/fake_firestore.py,
whose commit re-checks every read like Firestore does; ``before_commit`` holds
both transactions at the commit point so the interleaving actually happens.
"""
import asyncio
import threading
from datetime import datetime, timezone, timedelta
import pytest
from app.internal import discord_leases as leases
from app.internal import discord_streams as streams
from tests.fake_firestore import FakeFirestore
ORG = "org-a"
SYS = "sys-1"
@pytest.fixture
def fs(monkeypatch):
fake = FakeFirestore.install(monkeypatch)
published: list = []
def _publish(node_id, payload):
published.append((node_id, payload))
return fake.mqtt_up
fake.mqtt_up = True
fake.published = published
from app.internal.mqtt_handler import mqtt_handler
monkeypatch.setattr(mqtt_handler, "publish_discord_desired", _publish)
streams._publish_locks.clear()
async def _no_discord(token_ids):
return 0
monkeypatch.setattr(leases, "refresh_membership", _no_discord)
return fake
def _node(fake, node_id, *, org=ORG, system=SYS, status="online", age=0):
fake.put("nodes", node_id, {
"node_id": node_id, "name": node_id, "org_id": org, "status": status,
"approval_status": "approved", "assigned_system_id": system,
"last_seen": (datetime.now(timezone.utc) - timedelta(seconds=age)).isoformat(),
})
def _token(fake, token_id, *, org=None, guilds=None, valid=True, name=None):
fake.put("bot_tokens", token_id, {
"token_id": token_id, "name": name or token_id, "token": f"secret-{token_id}",
"org_id": org, "guild_ids": guilds, "valid": valid, "application_id": f"app-{token_id}",
"lease": None, "in_use": False, "assigned_node_id": None,
})
def _system(fake, system_id=SYS, org=ORG):
fake.put("systems", system_id, {"system_id": system_id, "name": "County P25", "org_id": org})
async def _start(guild="g1", channel="c1", **kw):
return await streams.start_stream(org_id=kw.pop("org_id", ORG), guild_id=guild, channel_id=channel,
system_id=kw.pop("system_id", SYS), **kw)
def _lease(fake, token_id):
return (fake.get("bot_tokens", token_id) or {}).get("lease")
# ---------------------------------------------------------------------------
# pick_token — pure choice
# ---------------------------------------------------------------------------
def test_pick_prefers_byo_then_reuse_and_never_crosses_orgs():
toks = {
"pool": {"name": "a", "guild_ids": ["g1"]},
"mine": {"name": "z", "org_id": ORG, "guild_ids": ["g1"]},
"theirs": {"name": "b", "org_id": "org-b", "guild_ids": ["g1"]},
}
assert leases.pick_token(toks, node_id="n1", guild_id="g1", org_id=ORG).token_id == "mine"
# Another org's branded bot is never offered, even when it is the only one.
res = leases.pick_token({"theirs": toks["theirs"]}, node_id="n1", guild_id="g1", org_id=ORG)
assert not res.ok and res.error == "pool_exhausted"
def test_pick_skips_other_nodes_busy_guilds_and_invalid_tokens():
toks = {
"other-node": {"lease": {"node_id": "n2", "streams": []}},
"busy-here": {"lease": {"node_id": "n1", "streams": ["g1:c9"]}},
"dead": {"valid": False},
}
res = leases.pick_token(toks, node_id="n1", guild_id="g1", org_id=ORG)
assert not res.ok and res.error == "pool_exhausted"
# The same session can serve a different guild (fan-out).
assert leases.pick_token(toks, node_id="n1", guild_id="g2", org_id=ORG).token_id == "busy-here"
def test_pick_reports_invite_links_when_no_bot_is_in_the_guild():
toks = {"t1": {"name": "Radio 1", "guild_ids": ["other"], "application_id": "123"}}
res = leases.pick_token(toks, node_id="n1", guild_id="g1", org_id=ORG)
assert res.error == "no_bot_in_guild"
assert res.invite[0]["name"] == "Radio 1"
assert "client_id=123" in res.invite[0]["invite_url"] and "guild_id=g1" in res.invite[0]["invite_url"]
assert leases.pick_token({}, node_id="n1", guild_id="g1", org_id=ORG).error == "pool_empty"
# ---------------------------------------------------------------------------
# Leasing through start_stream
# ---------------------------------------------------------------------------
async def test_rejoin_reuses_the_lease_instead_of_taking_a_second_token(fs):
_system(fs); _node(fs, "n1"); _token(fs, "t1"); _token(fs, "t2")
first = await _start()
again = await _start()
assert again["token_id"] == first["token_id"] and again["reused"] is True
leased = [t for t in ("t1", "t2") if _lease(fs, t)]
assert leased == [first["token_id"]] # server-26#204 item 2
assert _lease(fs, first["token_id"])["streams"] == ["g1:c1"]
async def test_preferred_token_is_honoured(fs):
_system(fs); _node(fs, "n1"); _token(fs, "t1"); _token(fs, "t2")
doc = await _start(preferred_token_id="t2")
assert doc["token_id"] == "t2" # server-26#204 item 1
async def test_one_session_fans_out_across_guilds_but_not_within_one(fs):
_system(fs); _node(fs, "n1"); _token(fs, "t1"); _token(fs, "t2")
a = await _start("g1", "c1")
b = await _start("g2", "c1")
c = await _start("g1", "c2")
assert a["token_id"] == b["token_id"] # one gateway session, two guilds
assert c["token_id"] != a["token_id"] # one voice connection per bot per guild
with pytest.raises(streams.StreamError) as e:
await _start("g1", "c3")
assert e.value.code == "pool_exhausted"
async def test_join_never_moves_another_guilds_stream(fs):
"""server-26#203 item 2: a /join from guild B used to steal guild A's node."""
_system(fs); _node(fs, "n1"); _token(fs, "t1")
await _start("gA", "c1")
await _start("gB", "c1")
a = fs.get("discord_streams", "gA:c1")
assert a["desired"] == "on" and a["token_id"] == "t1"
assert sorted(_lease(fs, "t1")["streams"]) == ["gA:c1", "gB:c1"]
async def test_desired_state_never_carries_the_token_string(fs):
_system(fs); _node(fs, "n1"); _token(fs, "t1")
await _start()
node_id, payload = fs.published[-1]
assert node_id == "n1" and payload["streams"][0]["token_id"] == "t1"
assert "secret-t1" not in repr(payload)
async def test_no_node_in_org_is_reported(fs):
_system(fs); _node(fs, "n1", org="org-b"); _token(fs, "t1")
with pytest.raises(streams.StreamError) as e:
await _start()
assert e.value.code == "no_node"
assert _lease(fs, "t1") is None
async def test_guild_stream_limit(fs):
_system(fs); _node(fs, "n1"); _token(fs, "t1"); _token(fs, "t2")
await _start("g1", "c1", max_streams=1)
with pytest.raises(streams.StreamError) as e:
await _start("g1", "c2", max_streams=1)
assert e.value.code == "limit"
# Re-requesting the channel that is already on is not a new stream.
await _start("g1", "c1", max_streams=1)
# ---------------------------------------------------------------------------
# Races — two concurrent requests can never both win (server-26#204 item 4)
# ---------------------------------------------------------------------------
def _hold_both_at_commit(fake):
barrier = threading.Barrier(2, timeout=5)
def _hook(attempt):
if attempt == 0:
barrier.wait()
fake.before_commit = _hook
async def test_two_nodes_racing_for_the_last_token(fs):
_node(fs, "n1"); _node(fs, "n2"); _token(fs, "t1")
_hold_both_at_commit(fs)
r1, r2 = await asyncio.gather(
leases.lease_for_stream("g1:c1", node_id="n1", org_id=ORG),
leases.lease_for_stream("g2:c1", node_id="n2", org_id=ORG),
)
assert fs.retries >= 1 # the race really happened
assert sorted([r1.ok, r2.ok]) == [False, True]
winner = "n1" if r1.ok else "n2"
loser_stream = "g2:c1" if r1.ok else "g1:c1"
lease = _lease(fs, "t1")
assert lease["node_id"] == winner and len(lease["streams"]) == 1
assert (fs.get("discord_streams", loser_stream) or {}).get("token_id") is None
async def test_same_channel_requested_twice_concurrently_takes_one_token(fs):
_node(fs, "n1"); _token(fs, "t1"); _token(fs, "t2")
_hold_both_at_commit(fs)
r1, r2 = await asyncio.gather(
leases.lease_for_stream("g1:c1", node_id="n1", org_id=ORG),
leases.lease_for_stream("g1:c1", node_id="n1", org_id=ORG),
)
assert r1.ok and r2.ok and r1.token_id == r2.token_id
held = [t for t in ("t1", "t2") if _lease(fs, t)]
assert held == [r1.token_id]
assert _lease(fs, r1.token_id)["streams"] == ["g1:c1"]
async def test_two_guilds_same_node_racing_share_one_session(fs):
_node(fs, "n1"); _token(fs, "t1")
_hold_both_at_commit(fs)
r1, r2 = await asyncio.gather(
leases.lease_for_stream("g1:c1", node_id="n1", org_id=ORG),
leases.lease_for_stream("g2:c1", node_id="n1", org_id=ORG),
)
assert r1.ok and r2.ok and r1.token_id == r2.token_id == "t1"
assert r1.lease_id == r2.lease_id
assert sorted(_lease(fs, "t1")["streams"]) == ["g1:c1", "g2:c1"]
# ---------------------------------------------------------------------------
# Ordered release (node-26#25): free only after the node confirms logout
# ---------------------------------------------------------------------------
async def test_stop_marks_releasing_and_holds_until_the_node_acks(fs):
_system(fs); _node(fs, "n1"); _node(fs, "n2"); _token(fs, "t1")
await _start()
await streams.stop_stream("g1:c1")
lease = _lease(fs, "t1")
assert lease["releasing"] is True and lease["streams"] == []
drop_seq = lease["release_seq"]
assert drop_seq == fs.published[-1][1]["seq"] # stamped with the list that dropped it
assert fs.published[-1][1]["streams"] == []
# Another node cannot take a token that may still be logged in on n1.
res = await leases.lease_for_stream("g9:c1", node_id="n2", org_id=ORG)
assert not res.ok
# A stale report (from before the node saw the drop) frees nothing.
assert await leases.settle_releases("n1", set(), drop_seq - 1) == []
# Up to date, but the bot is still logged in: nothing.
assert await leases.settle_releases("n1", {"t1"}, drop_seq) == []
# Up to date and logged out: freed.
assert await leases.settle_releases("n1", set(), drop_seq) == ["t1"]
tok = fs.get("bot_tokens", "t1")
assert tok["lease"] is None and tok["in_use"] is False and tok["assigned_node_id"] is None
async def test_restart_before_release_reuses_the_logged_in_session(fs):
_system(fs); _node(fs, "n1"); _token(fs, "t1"); _token(fs, "t2")
first = await _start()
await streams.stop_stream("g1:c1")
again = await _start()
assert again["token_id"] == first["token_id"]
lease = _lease(fs, first["token_id"])
assert lease["releasing"] is False and lease["streams"] == ["g1:c1"]
async def test_handle_actual_settles_and_ignores_other_nodes_streams(fs):
_system(fs); _node(fs, "n1"); _node(fs, "n2"); _token(fs, "t1"); _token(fs, "t2")
await _start("g1", "c1", node_id="n1")
await _start("g2", "c1", node_id="n2")
await streams.handle_actual("n1", {
"applied_seq": 1, "sessions": ["t1"],
"streams": {"g1:c1": {"state": "connected"}, "g2:c1": {"state": "error", "error": "spoof"}},
})
assert fs.get("discord_streams", "g1:c1")["actual"]["state"] == "connected"
assert fs.get("discord_streams", "g2:c1")["actual"]["state"] == "pending" # not n1's to report
await streams.stop_stream("g1:c1")
seq = _lease(fs, "t1")["release_seq"]
await streams.handle_actual("n1", {"applied_seq": seq, "sessions": [], "streams": {}})
assert _lease(fs, "t1") is None
async def test_failed_publish_is_retried_by_the_reconciler(fs):
"""server-26#204 item 3: a send failure used to leak the token forever."""
_system(fs); _node(fs, "n1"); _token(fs, "t1")
fs.mqtt_up = False
doc = await _start()
assert doc["delivered"] is False
assert fs.get("discord_nodes", "n1")["published"] is False
fs.mqtt_up = True
n = len(fs.published)
await streams.reconcile_once()
assert len(fs.published) > n and fs.get("discord_nodes", "n1")["published"] is True
assert fs.published[-1][1]["streams"][0]["stream_id"] == "g1:c1"
async def test_concurrent_publishes_keep_the_retained_list_newest(fs):
_system(fs); _node(fs, "n1"); _token(fs, "t1")
await _start()
await asyncio.gather(*(streams.publish_desired("n1") for _ in range(5)))
seqs = [p["seq"] for nid, p in fs.published if nid == "n1"]
assert seqs == sorted(seqs) and len(set(seqs)) == len(seqs)
# ---------------------------------------------------------------------------
# Guild scoping (server-26#203 item 1)
# ---------------------------------------------------------------------------
async def test_leave_in_one_guild_never_touches_another(fs):
_system(fs); _node(fs, "n1"); _token(fs, "t1"); _token(fs, "t2")
await _start("gA", "c1")
await _start("gB", "c1")
stopped = await streams.stop_guild("gA", "leave")
assert stopped == ["gA:c1"]
assert fs.get("discord_streams", "gB:c1")["desired"] == "on"
assert fs.get("discord_streams", "gA:c1")["desired"] == "off"
async def test_guild_org_resolution(fs, monkeypatch):
monkeypatch.setattr(streams.settings, "discord_unlinked_guild_org_id", "founding")
fs.put("discord_guilds", "g-linked", {"org_id": "org-b"})
assert await streams.resolve_guild_org("g-linked") == "org-b"
assert await streams.resolve_guild_org("g-other") == "founding"
monkeypatch.setattr(streams.settings, "discord_unlinked_guild_org_id", "")
assert await streams.resolve_guild_org("g-other") is None
# ---------------------------------------------------------------------------
# Reconciler — offline grace, re-placement
# ---------------------------------------------------------------------------
async def test_offline_node_keeps_its_lease_inside_the_grace(fs, monkeypatch):
monkeypatch.setattr(streams.settings, "discord_lease_offline_release_seconds", 600)
_system(fs); _node(fs, "n1"); _token(fs, "t1")
await _start(mode="persistent")
_node(fs, "n1", status="offline", age=120)
await streams.reconcile_once()
assert _lease(fs, "t1")["node_id"] == "n1"
assert fs.get("discord_streams", "g1:c1")["actual"]["state"] == "node_offline"
async def test_persistent_stream_moves_after_grace_adhoc_stops(fs, monkeypatch):
monkeypatch.setattr(streams.settings, "discord_lease_offline_release_seconds", 600)
_system(fs); _node(fs, "n1"); _token(fs, "t1"); _token(fs, "t2")
await _start("g1", "c1", mode="persistent")
await _start("g2", "c1", mode="adhoc")
_node(fs, "n1", status="offline", age=900)
_node(fs, "n2")
await streams.reconcile_once()
moved = fs.get("discord_streams", "g1:c1")
assert moved["desired"] == "on" and moved["node_id"] == "n2" and moved["token_id"]
assert _lease(fs, moved["token_id"])["node_id"] == "n2"
assert fs.get("discord_streams", "g2:c1")["desired"] == "off"
assert all((_lease(fs, t) or {}).get("node_id") != "n1" for t in ("t1", "t2"))
async def test_leave_while_node_offline_frees_token_after_grace(fs, monkeypatch):
"""A stream stopped while its node was down leaves no desired=on stream
behind; the token must still return to the pool once the grace passes."""
monkeypatch.setattr(streams.settings, "discord_lease_offline_release_seconds", 600)
_system(fs); _node(fs, "n1"); _token(fs, "t1")
await _start()
_node(fs, "n1", status="offline", age=120)
await streams.stop_stream("g1:c1")
await streams.reconcile_once()
assert _lease(fs, "t1")["releasing"] is True # inside the grace: held
_node(fs, "n1", status="offline", age=900)
await streams.reconcile_once()
assert _lease(fs, "t1") is None
async def test_deleted_node_lease_freed_after_grace_from_last_report(fs, monkeypatch):
monkeypatch.setattr(streams.settings, "discord_lease_offline_release_seconds", 600)
_system(fs); _node(fs, "n1"); _token(fs, "t1")
await _start()
await streams.stop_stream("g1:c1")
fs.put("discord_nodes", "n1", {**(fs.get("discord_nodes", "n1") or {}),
"actual_at": datetime.now(timezone.utc).isoformat()})
await fs.doc_delete("nodes", "n1")
await streams.reconcile_once()
assert _lease(fs, "t1") is not None # heard from recently: held
old = (datetime.now(timezone.utc) - timedelta(seconds=900)).isoformat()
fs.put("discord_nodes", "n1", {**fs.get("discord_nodes", "n1"), "actual_at": old})
await streams.reconcile_once()
assert _lease(fs, "t1") is None
async def test_release_seq_never_names_a_list_that_still_holds_the_stream(fs, monkeypatch):
"""A stop landing between building the desired list and stamping must not
be stamped with that list's seq (the list still contains the stream)."""
_system(fs); _node(fs, "n1"); _token(fs, "t1")
await _start()
real = streams.desired_for_node
fired = []
async def _racy(node_id):
out = await real(node_id)
if not fired:
fired.append(1)
await leases.detach_stream("g1:c1", {"desired": "off"})
return out
monkeypatch.setattr(streams, "desired_for_node", _racy)
await streams.publish_desired("n1")
racy_seq = fs.published[-1][1]["seq"]
assert any(d["stream_id"] == "g1:c1" for d in fs.published[-1][1]["streams"])
rseq = _lease(fs, "t1")["release_seq"]
assert rseq is None or rseq > racy_seq
assert await leases.settle_releases("n1", set(), racy_seq) == []
async def test_retuned_node_gives_up_a_persistent_stream(fs):
_system(fs); _system(fs, "sys-2"); _node(fs, "n1"); _token(fs, "t1")
await _start(mode="persistent")
_node(fs, "n1", system="sys-2")
await streams.reconcile_once()
s = fs.get("discord_streams", "g1:c1")
assert s["desired"] == "on" and s["token_id"] is None and s["actual"]["state"] == "error"
assert _lease(fs, "t1")["releasing"] is True
# ---------------------------------------------------------------------------
# Admin force release + token pickup endpoint
# ---------------------------------------------------------------------------
async def test_force_release_detaches_streams(fs):
_system(fs); _node(fs, "n1"); _token(fs, "t1")
await _start()
old = await leases.force_release_token("t1")
assert old["node_id"] == "n1"
assert _lease(fs, "t1") is None
assert fs.get("discord_streams", "g1:c1")["token_id"] is None
def test_node_token_pickup_only_for_the_lease_holder(fs):
from fastapi import FastAPI
from fastapi.testclient import TestClient
from app.routers import discord as discord_router
fs.put("node_keys", "n1", {"node_id": "n1", "api_key": "key-1"})
fs.put("node_keys", "n2", {"node_id": "n2", "api_key": "key-2"})
_token(fs, "t1")
fs.put("bot_tokens", "t1", {**fs.get("bot_tokens", "t1"),
"lease": {"node_id": "n1", "lease_id": "L", "streams": ["g:c"], "releasing": False}})
app = FastAPI()
app.include_router(discord_router.router)
client = TestClient(app)
ok = client.get("/discord/node/n1/tokens/t1", headers={"Authorization": "Bearer key-1"})
assert ok.status_code == 200 and ok.json()["token"] == "secret-t1"
assert client.get("/discord/node/n2/tokens/t1", headers={"Authorization": "Bearer key-2"}).status_code == 403
assert client.get("/discord/node/n1/tokens/t1", headers={"Authorization": "Bearer key-2"}).status_code == 401
assert client.get("/discord/node/n1/tokens/t1").status_code == 401
# ---------------------------------------------------------------------------
# Wiring
# ---------------------------------------------------------------------------
async def test_mqtt_dispatch_routes_discord_actual(monkeypatch):
from app.internal.mqtt_handler import MQTTHandler
seen = []
async def _fake(node_id, payload):
seen.append((node_id, payload))
monkeypatch.setattr(streams, "handle_actual", _fake)
await MQTTHandler()._dispatch("nodes/n1/discord/actual", {"applied_seq": 3})
assert seen == [("n1", {"applied_seq": 3})]
def test_command_payload_keeps_preferred_token_id():
from app.models import CommandPayload
cmd = CommandPayload(action="discord_join", guild_id="g", channel_id="c", preferred_token_id="t9")
assert cmd.preferred_token_id == "t9"
Reference in New Issue View Git Blame Copy Permalink