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

@@ -71,6 +71,59 @@ class StateStoreCommitTest(unittest.TestCase):
             self.assertEqual(list(root.glob(".*.tmp")), [])
 
 
+class StateStoreValidationTest(unittest.TestCase):
+    @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 test_impossible_or_over_capacity_token_state_is_rejected(self) -> None:
+        session_id = "1" * 64
+        session = {
+            "anchor_pid": 123,
+            "anchor_starttime": "456",
+            "runtime_generation": 2,
+        }
+        live_token = {
+            "session_id": session_id,
+            "runtime_generation": 2,
+            "binding": self.binding(),
+            "consumed": False,
+        }
+        cases = {
+            "stale generation": {
+                "2" * 64: {**live_token, "runtime_generation": 1},
+            },
+            "consumed token": {
+                "2" * 64: {**live_token, "consumed": True},
+            },
+            "over capacity": {
+                f"{index:064x}": copy.deepcopy(live_token)
+                for index in range(DAEMON.MAX_PENDING_TOKENS + 1)
+            },
+        }
+
+        for label, tokens in cases.items():
+            with self.subTest(label=label), tempfile.TemporaryDirectory() as directory:
+                root = Path(directory)
+                os.chmod(root, 0o700)
+                state_path = root / "state.json"
+                state_path.write_text(json.dumps({
+                    "version": 1,
+                    "sessions": {session_id: session},
+                    "tokens": tokens,
+                }))
+                os.chmod(state_path, 0o600)
+
+                with self.assertRaisesRegex(DAEMON.BrokerFailure, "STATE_INTEGRITY"):
+                    DAEMON.StateStore(state_path)
+
+
 class BrokerBehaviorTest(unittest.TestCase):
     def make_broker(self, root: Path):
         os.chmod(root, 0o700)
@@ -204,7 +257,41 @@ class BrokerBehaviorTest(unittest.TestCase):
             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:
+    def test_directory_fsync_failure_poisoned_store_cannot_continue(self) -> None:
+        with tempfile.TemporaryDirectory() as directory, patch.object(
+            DAEMON,
+            "proc_node",
+            return_value={"pid": 123, "ppid": 1, "starttime": "456"},
+        ):
+            root = Path(directory)
+            broker = self.make_broker(root)
+            real_fsync = os.fsync
+
+            def fail_directory_fsync(descriptor: int) -> None:
+                if os.path.isdir(f"/proc/self/fd/{descriptor}"):
+                    raise OSError("directory fsync failed")
+                real_fsync(descriptor)
+
+            with patch.object(DAEMON.os, "fsync", side_effect=fail_directory_fsync):
+                with self.assertRaisesRegex(
+                    DAEMON.StateCommitUncertain, "STATE_COMMIT_UNCERTAIN"
+                ):
+                    self.register(broker)
+
+            durable = json.loads(broker.store.path.read_text())
+            self.assertEqual(broker.store.value, durable)
+            self.assertTrue(broker.store.poisoned)
+            before = copy.deepcopy(broker.store.value)
+            with self.assertRaisesRegex(
+                DAEMON.StateCommitUncertain, "STATE_COMMIT_UNCERTAIN"
+            ):
+                broker.handle((123, 1000, 1000), {
+                    "action": "register_anchor",
+                    "runtime_generation": 2,
+                })
+            self.assertEqual(broker.store.value, before)
+
+    def test_commit_failures_before_replace_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):