diff --git a/gates/gate_bus.py b/gates/gate_bus.py index 4013ed1..fd735e2 100644 --- a/gates/gate_bus.py +++ b/gates/gate_bus.py @@ -9,12 +9,14 @@ class GateCancelled(Exception): class GateBus: messages = {} # node_id(str) -> chosen int (1-based) masks = {} # node_id(str) -> PNG bytes + payloads = {} # node_id(str) -> arbitrary payload (e.g., edited text) cancelled = False @classmethod def arm(cls, node_id): cls.messages.pop(str(node_id), None) cls.masks.pop(str(node_id), None) + cls.payloads.pop(str(node_id), None) cls.cancelled = False @classmethod @@ -41,3 +43,17 @@ class GateBus: @classmethod def pop_mask(cls, node_id): return cls.masks.pop(str(node_id), None) + + @classmethod + def put_payload(cls, node_id, value): + cls.payloads[str(node_id)] = value + + @classmethod + def wait_payload(cls, node_id, period=0.1, should_cancel=None): + sid = str(node_id) + while sid not in cls.payloads: + if cls.cancelled or (should_cancel is not None and should_cancel()): + cls.cancelled = False + raise GateCancelled() + time.sleep(period) + return cls.payloads.pop(sid) diff --git a/tests/test_gate_bus.py b/tests/test_gate_bus.py index 231b29b..0cd9962 100644 --- a/tests/test_gate_bus.py +++ b/tests/test_gate_bus.py @@ -36,3 +36,26 @@ def test_arm_clears_mask(): gb.GateBus.put_mask("9", b"x") gb.GateBus.arm("9") assert gb.GateBus.pop_mask("9") is None + +def test_payload_roundtrip(): + gb.GateBus.arm("p") + gb.GateBus.put_payload("p", "hello edited") + assert gb.GateBus.wait_payload("p") == "hello edited" + +def test_payload_consumed(): + gb.GateBus.arm("p") + gb.GateBus.put_payload("p", "x") + gb.GateBus.wait_payload("p") + assert "p" not in gb.GateBus.payloads + +def test_arm_clears_payload(): + gb.GateBus.put_payload("p", "stale") + gb.GateBus.arm("p") + assert "p" not in gb.GateBus.payloads + +def test_wait_payload_cancel_flag_raises(): + import pytest + gb.GateBus.arm("p") + gb.GateBus.cancelled = True + with pytest.raises(gb.GateCancelled): + gb.GateBus.wait_payload("p")