Files
stack/packages/mosaic/src/lease-broker/state_store_unittest.py
jason.woltje 8ec67a1126
ci/woodpecker/push/publish Pipeline was successful
ci/woodpecker/push/ci Pipeline was successful
feat(mosaic): add authenticated external lease broker (#836)
2026-07-18 03:12:23 +00:00

334 lines
14 KiB
Python

#!/usr/bin/env python3
"""Standard-library edge tests for lease-broker atomic state persistence."""
from __future__ import annotations
import copy
import importlib.util
import json
import os
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch
DAEMON_PATH = Path(__file__).parents[2] / "framework/tools/lease-broker/daemon.py"
SPEC = importlib.util.spec_from_file_location("lease_broker_daemon", DAEMON_PATH)
if SPEC is None or SPEC.loader is None:
raise RuntimeError("unable to load lease broker daemon")
DAEMON = importlib.util.module_from_spec(SPEC)
SPEC.loader.exec_module(DAEMON)
class StateStoreCommitTest(unittest.TestCase):
def make_store(self, root: Path):
os.chmod(root, 0o700)
store = DAEMON.StateStore(root / "state.json")
store.value["marker"] = "partial-write-proof"
return store
def test_partial_writes_persist_the_complete_payload(self) -> None:
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
store = self.make_store(root)
real_write = os.write
def partial_write(descriptor: int, payload: bytes) -> int:
return real_write(descriptor, payload[: max(1, len(payload) // 3)])
with patch.object(DAEMON.os, "write", side_effect=partial_write):
store.commit()
self.assertEqual(json.loads(store.path.read_text()), store.value)
def test_zero_progress_removes_owned_temporary_file(self) -> None:
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
store = self.make_store(root)
with patch.object(DAEMON.os, "write", return_value=0):
with self.assertRaises(OSError):
store.commit()
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 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)
return DAEMON.Broker(DAEMON.StateStore(root / "state.json"))
@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:
root = Path(directory)
broker = self.make_broker(root)
with (
patch.object(
DAEMON,
"proc_node",
return_value={"pid": 123, "ppid": 1, "starttime": "456"},
),
patch.object(DAEMON, "verified_ancestry", return_value=True),
):
first = broker.handle((123, 1000, 1000), {
"action": "register_anchor",
"runtime_generation": 1,
})
minted = broker.handle((123, 1000, 1000), {
"action": "mint_token",
"session_id": first["session_id"],
"runtime_generation": 1,
"binding": self.binding(),
})
bumped = broker.handle((123, 1000, 1000), {
"action": "register_anchor",
"runtime_generation": 2,
})
repeated = broker.handle((123, 1000, 1000), {
"action": "register_anchor",
"runtime_generation": 2,
})
self.assertEqual(bumped["session_id"], first["session_id"])
self.assertEqual(repeated["session_id"], first["session_id"])
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), {
"action": "register_anchor",
"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_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):
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()