test(#828): cover bounded broker token state

packages/mosaic/src/lease-broker/state_store_unittest.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import copy
 import importlib.util
 import json
 import os
@@ -52,22 +53,67 @@ class StateStoreCommitTest(unittest.TestCase):
             self.assertFalse(store.path.exists())
             self.assertEqual(list(root.glob(".*.tmp")), [])
 
+    def test_oversized_payload_is_refused_before_replacing_state(self) -> None:
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            store = self.make_store(root)
+            store.value.pop("marker")
+            store.commit()
+            durable = store.path.read_bytes()
+            store.value["oversized"] = "x" * DAEMON.MAX_STATE
+
+            with patch.object(DAEMON.os, "open", wraps=os.open) as mocked_open:
+                with self.assertRaisesRegex(DAEMON.BrokerFailure, "STATE_TOO_LARGE"):
+                    store.commit()
+
+            self.assertEqual(mocked_open.call_count, 0)
+            self.assertEqual(store.path.read_bytes(), durable)
+            self.assertEqual(list(root.glob(".*.tmp")), [])
+
 
 class BrokerBehaviorTest(unittest.TestCase):
     def make_broker(self, root: Path):
         os.chmod(root, 0o700)
         return DAEMON.Broker(DAEMON.StateStore(root / "state.json"))
 
-    def test_anchor_generation_bump_reuses_session_and_revokes_token(self) -> None:
-        binding = {
+    @staticmethod
+    def binding() -> dict[str, object]:
+        return {
             "compaction_epoch": 0,
             "request_epoch": 0,
             "h_source": "a" * 64,
             "h_payload": "b" * 64,
             "schema_version": 1,
         }
+
+    def register(self, broker, generation: int = 1) -> str:
+        response = broker.handle((123, 1000, 1000), {
+            "action": "register_anchor",
+            "runtime_generation": generation,
+        })
+        return response["session_id"]
+
+    def mint(self, broker, session_id: str, generation: int = 1) -> str:
+        response = broker.handle((123, 1000, 1000), {
+            "action": "mint_token",
+            "session_id": session_id,
+            "runtime_generation": generation,
+            "binding": self.binding(),
+        })
+        return response["token"]
+
+    def consume(self, broker, session_id: str, token: str, generation: int = 1):
+        return broker.handle((123, 1000, 1000), {
+            "action": "consume_token",
+            "session_id": session_id,
+            "runtime_generation": generation,
+            "token": token,
+        })
+
+    def test_anchor_generation_bump_reuses_session_and_revokes_token(self) -> None:
         with tempfile.TemporaryDirectory() as directory:
-            broker = self.make_broker(Path(directory))
+            root = Path(directory)
+            broker = self.make_broker(root)
             with (
                 patch.object(
                     DAEMON,
@@ -84,7 +130,7 @@ class BrokerBehaviorTest(unittest.TestCase):
                     "action": "mint_token",
                     "session_id": first["session_id"],
                     "runtime_generation": 1,
-                    "binding": binding,
+                    "binding": self.binding(),
                 })
                 bumped = broker.handle((123, 1000, 1000), {
                     "action": "register_anchor",
@@ -97,7 +143,12 @@ class BrokerBehaviorTest(unittest.TestCase):
 
             self.assertEqual(bumped["session_id"], first["session_id"])
             self.assertEqual(repeated["session_id"], first["session_id"])
-            self.assertTrue(broker.store.tokens()[minted["token"]]["consumed"])
+            self.assertNotIn(minted["token"], broker.store.tokens())
+            restarted = self.make_broker(root)
+            self.assertEqual(restarted.store.tokens(), {})
+            self.assertEqual(
+                restarted.store.sessions()[first["session_id"]]["runtime_generation"], 2
+            )
             with self.assertRaisesRegex(DAEMON.BrokerFailure, "STALE_GENERATION"):
                 with patch.object(DAEMON, "proc_node", return_value={"starttime": "456"}):
                     broker.handle((123, 1000, 1000), {
@@ -105,5 +156,91 @@ class BrokerBehaviorTest(unittest.TestCase):
                         "runtime_generation": 1,
                     })
 
+    def test_successful_consume_deletes_token_and_replay_is_refused(self) -> None:
+        with tempfile.TemporaryDirectory() as directory, (
+            patch.object(DAEMON, "proc_node", return_value={"pid": 123, "ppid": 1, "starttime": "456"})
+        ), patch.object(DAEMON, "verified_ancestry", return_value=True):
+            broker = self.make_broker(Path(directory))
+            session_id = self.register(broker)
+            token = self.mint(broker, session_id)
+
+            self.assertEqual(self.consume(broker, session_id, token), {"ok": True})
+            self.assertNotIn(token, broker.store.tokens())
+            with self.assertRaisesRegex(DAEMON.BrokerFailure, "TOKEN_REPLAY"):
+                self.consume(broker, session_id, token)
+
+    def test_normal_cycles_remain_bounded_and_restartable(self) -> None:
+        with tempfile.TemporaryDirectory() as directory, (
+            patch.object(DAEMON, "proc_node", return_value={"pid": 123, "ppid": 1, "starttime": "456"})
+        ), patch.object(DAEMON, "verified_ancestry", return_value=True):
+            root = Path(directory)
+            broker = self.make_broker(root)
+            session_id = self.register(broker)
+
+            for _ in range(DAEMON.MAX_PENDING_TOKENS * 3):
+                self.consume(broker, session_id, self.mint(broker, session_id))
+
+            self.assertEqual(broker.store.tokens(), {})
+            self.assertLess((root / "state.json").stat().st_size, DAEMON.MAX_STATE)
+            restarted = self.make_broker(root)
+            self.assertEqual(restarted.store.tokens(), {})
+            self.assertIn(session_id, restarted.store.sessions())
+
+    def test_pending_token_capacity_refusal_does_not_mutate_state(self) -> None:
+        with tempfile.TemporaryDirectory() as directory, (
+            patch.object(DAEMON, "proc_node", return_value={"pid": 123, "ppid": 1, "starttime": "456"})
+        ), patch.object(DAEMON, "verified_ancestry", return_value=True):
+            root = Path(directory)
+            broker = self.make_broker(root)
+            session_id = self.register(broker)
+            for _ in range(DAEMON.MAX_PENDING_TOKENS):
+                self.mint(broker, session_id)
+            before = copy.deepcopy(broker.store.value)
+            durable = broker.store.path.read_bytes()
+
+            with self.assertRaisesRegex(DAEMON.BrokerFailure, "TOKEN_CAPACITY"):
+                self.mint(broker, session_id)
+
+            self.assertEqual(broker.store.value, before)
+            self.assertEqual(broker.store.path.read_bytes(), durable)
+
+    def test_commit_failures_roll_back_every_broker_mutation(self) -> None:
+        with tempfile.TemporaryDirectory() as directory, (
+            patch.object(DAEMON, "proc_node", return_value={"pid": 123, "ppid": 1, "starttime": "456"})
+        ), patch.object(DAEMON, "verified_ancestry", return_value=True):
+            root = Path(directory)
+            broker = self.make_broker(root)
+            session_id = self.register(broker)
+            token = self.mint(broker, session_id)
+
+            def assert_rollback(request: dict[str, object]) -> None:
+                before = copy.deepcopy(broker.store.value)
+                durable = broker.store.path.read_bytes()
+                with patch.object(broker.store, "commit", side_effect=OSError("fsync failed")):
+                    with self.assertRaisesRegex(OSError, "fsync failed"):
+                        broker.handle((123, 1000, 1000), request)
+                self.assertEqual(broker.store.value, before)
+                self.assertEqual(broker.store.path.read_bytes(), durable)
+
+            assert_rollback({"action": "register_anchor", "runtime_generation": 2})
+            assert_rollback({
+                "action": "mint_token", "session_id": session_id,
+                "runtime_generation": 1, "binding": self.binding(),
+            })
+            assert_rollback({
+                "action": "consume_token", "session_id": session_id,
+                "runtime_generation": 1, "token": token,
+            })
+
+            with tempfile.TemporaryDirectory() as second_directory:
+                second = self.make_broker(Path(second_directory))
+                with patch.object(second.store, "commit", side_effect=OSError("fsync failed")):
+                    with self.assertRaisesRegex(OSError, "fsync failed"):
+                        self.register(second)
+                self.assertEqual(
+                    second.store.value, {"version": 1, "sessions": {}, "tokens": {}}
+                )
+                self.assertFalse(second.store.path.exists())
+
 if __name__ == "__main__":
     unittest.main()