diff --git a/packages/mosaic/src/lease-broker/state_store_unittest.py b/packages/mosaic/src/lease-broker/state_store_unittest.py index bceb696..5549570 100644 --- a/packages/mosaic/src/lease-broker/state_store_unittest.py +++ b/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()