#!/usr/bin/env python3 """Branch-focused tests for the lease-gated runtime executables.""" from __future__ import annotations import importlib.util import io import json import os import runpy import socket import stat import subprocess import sys import tempfile import threading import unittest from contextlib import redirect_stderr from pathlib import Path from types import SimpleNamespace from unittest.mock import patch TOOLS_DIR = Path(__file__).parents[2] / "framework/tools/lease-broker" if str(TOOLS_DIR) not in sys.path: sys.path.insert(0, str(TOOLS_DIR)) def load_tool(module_name: str, filename: str): spec = importlib.util.spec_from_file_location(module_name, TOOLS_DIR / filename) if spec is None or spec.loader is None: raise RuntimeError(f"unable to load {filename}") module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) return module LAUNCHER = load_tool("lease_runtime_launcher", "launch-runtime.py") GATE = load_tool("lease_mutator_gate", "mutator-gate.py") class FakeSocket: def __init__(self, *chunks: bytes): self.chunks = list(chunks) self.timeout = None self.connected = None self.sent = b"" self.shutdown_how = None def __enter__(self): return self def __exit__(self, *_args): return False def settimeout(self, value: float) -> None: self.timeout = value def connect(self, value: str) -> None: self.connected = value def sendall(self, value: bytes) -> None: self.sent += value def shutdown(self, how: int) -> None: self.shutdown_how = how def recv(self, _size: int) -> bytes: return self.chunks.pop(0) if self.chunks else b"" class LaunchRuntimeTest(unittest.TestCase): def test_success_registers_then_injects_session_before_exec(self) -> None: calls: dict[str, object] = {} session_id = "a" * 64 def request(path: Path, payload: dict[str, object]) -> dict[str, object]: calls["path"] = path calls["request"] = payload return {"ok": True, "session_id": session_id} def execute(command: str, argv: list[str], environment: dict[str, str]) -> None: calls["execute"] = (command, argv, environment) def initialize_generation(path: Path, generation: int) -> None: calls["generation"] = (path, generation) result = LAUNCHER.main( ["--runtime", "claude", "--", "claude", "--print", "hello"], environ={ "MOSAIC_LEASE_BROKER_SOCKET": "/run/test/broker.sock", "MOSAIC_RUNTIME_GENERATION": "7", "PRESERVED": "yes", }, request=request, execute=execute, initialize_generation=initialize_generation, ) self.assertEqual(result, 0) self.assertEqual(calls["path"], Path("/run/test/broker.sock")) self.assertEqual( calls["request"], {"action": "register_anchor", "runtime_generation": 7}, ) command, argv, environment = calls["execute"] self.assertEqual(command, "claude") self.assertEqual(argv, ["claude", "--print", "hello"]) self.assertEqual(environment["MOSAIC_LEASE_SESSION_ID"], session_id) self.assertEqual(environment["MOSAIC_RUNTIME_GENERATION"], "7") self.assertEqual(environment["MOSAIC_LEASE_RUNTIME"], "claude") self.assertEqual( environment["MOSAIC_LEASE_GENERATION_FILE"], f"/run/test/generation-{session_id}.state", ) self.assertEqual( calls["generation"], (Path(f"/run/test/generation-{session_id}.state"), 7), ) self.assertEqual(environment["PRESERVED"], "yes") def test_dangerous_claude_mode_is_owned_and_injected_by_the_wrapper(self) -> None: executed: list[tuple[str, list[str], dict[str, str]]] = [] result = LAUNCHER.main( ["--runtime", "claude", "--dangerous", "--", "claude", "-p", "hello"], environ={"MOSAIC_LEASE_BROKER_SOCKET": "/broker"}, request=lambda *_args: {"ok": True, "session_id": "e" * 64}, execute=lambda *args: executed.append(args), initialize_generation=lambda *_args: None, ) self.assertEqual(result, 0) self.assertEqual( executed[0][1], ["claude", "--dangerously-skip-permissions", "-p", "hello"], ) with redirect_stderr(io.StringIO()): self.assertEqual( LAUNCHER.main( ["--runtime", "pi", "--dangerous", "--", "pi"], environ={"MOSAIC_LEASE_BROKER_SOCKET": "/broker"}, request=lambda *_args: {"ok": True, "session_id": "e" * 64}, execute=lambda *_args: self.fail("invalid dangerous runtime executed"), ), 64, ) def test_command_without_separator_is_forwarded_unchanged(self) -> None: executed: list[tuple[str, list[str], dict[str, str]]] = [] result = LAUNCHER.main( ["--runtime", "pi", "pi", "--print", "hello"], environ={"MOSAIC_LEASE_BROKER_SOCKET": "/broker"}, request=lambda *_args: {"ok": True, "session_id": "f" * 64}, execute=lambda *args: executed.append(args), initialize_generation=lambda *_args: None, ) self.assertEqual(result, 0) self.assertEqual(executed[0][0:2], ("pi", ["pi", "--print", "hello"])) def test_missing_command_is_usage_error(self) -> None: with redirect_stderr(io.StringIO()): self.assertEqual( LAUNCHER.main( ["--runtime", "pi", "--"], environ={}, request=lambda *_args: {}, execute=lambda *_args: None, ), 64, ) def test_registration_validation_and_environment_fail_closed(self) -> None: good_session = "b" * 64 cases = [ ({}, {"ok": True, "session_id": good_session}), ({"MOSAIC_LEASE_BROKER_SOCKET": "/x", "MOSAIC_RUNTIME_GENERATION": "bad"}, {}), ({"MOSAIC_LEASE_BROKER_SOCKET": "/x", "MOSAIC_RUNTIME_GENERATION": "-1"}, {}), ({"MOSAIC_LEASE_BROKER_SOCKET": "/x"}, {"ok": False, "session_id": good_session}), ({"MOSAIC_LEASE_BROKER_SOCKET": "/x"}, {"ok": True, "session_id": 4}), ({"MOSAIC_LEASE_BROKER_SOCKET": "/x"}, {"ok": True, "session_id": "b" * 63}), ({"MOSAIC_LEASE_BROKER_SOCKET": "/x"}, {"ok": True, "session_id": "z" * 64}), ] for environment, reply in cases: with self.subTest(environment=environment, reply=reply), redirect_stderr(io.StringIO()): executed: list[object] = [] result = LAUNCHER.main( ["--runtime", "pi", "--", "pi"], environ=environment, request=lambda *_args, value=reply: value, execute=lambda *args: executed.append(args), ) self.assertEqual(result, 1) self.assertEqual(executed, []) def test_generation_initialization_failure_denies_before_exec(self) -> None: with redirect_stderr(io.StringIO()): self.assertEqual( LAUNCHER.main( ["--runtime", "pi", "--", "pi"], environ={"MOSAIC_LEASE_BROKER_SOCKET": "/broker"}, request=lambda *_args: {"ok": True, "session_id": "a" * 64}, execute=lambda *_args: self.fail("must not execute"), initialize_generation=lambda *_args: (_ for _ in ()).throw( OSError("unsafe state") ), ), 1, ) def test_registration_exceptions_fail_closed(self) -> None: failures = [ValueError("bad"), OSError("down"), json.JSONDecodeError("bad", "x", 0)] for failure in failures: with self.subTest(failure=type(failure).__name__), redirect_stderr(io.StringIO()): def request(*_args, error=failure): raise error self.assertEqual( LAUNCHER.main( ["--runtime", "claude", "--", "claude"], environ={"MOSAIC_LEASE_BROKER_SOCKET": "/x"}, request=request, execute=lambda *_args: self.fail("must not execute"), ), 1, ) def test_exec_failure_is_fail_closed(self) -> None: with redirect_stderr(io.StringIO()): self.assertEqual( LAUNCHER.main( ["--runtime", "pi", "--", "pi"], environ={"MOSAIC_LEASE_BROKER_SOCKET": "/x"}, request=lambda *_args: {"ok": True, "session_id": "c" * 64}, execute=lambda *_args: (_ for _ in ()).throw(OSError("missing")), initialize_generation=lambda *_args: None, ), 1, ) def test_broker_reply_framing_and_shape_validation(self) -> None: replies = [ (b'{"ok":true}\n', {"ok": True}), (b'{"ok":true}', ValueError), (b'[]\n', ValueError), (b"x" * (LAUNCHER.MAX_FRAME + 1), ValueError), ] for wire_reply, expected in replies: with self.subTest(size=len(wire_reply)): fake = FakeSocket(wire_reply) with patch.object(LAUNCHER.socket, "socket", return_value=fake): if isinstance(expected, type) and issubclass(expected, Exception): with self.assertRaises(expected): LAUNCHER.broker_request(Path("/broker"), {"action": "register_anchor"}) else: self.assertEqual( LAUNCHER.broker_request(Path("/broker"), {"action": "register_anchor"}), expected, ) self.assertEqual(fake.timeout, LAUNCHER.BROKER_TIMEOUT_SECONDS) self.assertEqual(fake.connected, "/broker") self.assertEqual(fake.shutdown_how, socket.SHUT_WR) class ExecutableEntrypointTest(unittest.TestCase): def test_real_claude_and_pi_gates_fail_closed_on_empty_or_truncated_reply(self) -> None: for runtime in ("claude", "pi"): for wire_reply in (b"", b'{"ok":true'): with self.subTest(runtime=runtime, wire_reply=wire_reply), tempfile.TemporaryDirectory() as root: socket_path = Path(root) / "broker.sock" server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) server.bind(str(socket_path)) server.listen(1) def serve_reply() -> None: with server: connection, _ = server.accept() with connection: while connection.recv(4096): pass if wire_reply: connection.sendall(wire_reply) thread = threading.Thread(target=serve_reply, daemon=True) thread.start() environment = { **os.environ, "MOSAIC_LEASE_BROKER_SOCKET": str(socket_path), "MOSAIC_LEASE_SESSION_ID": "d" * 64, "MOSAIC_RUNTIME_GENERATION": "1", } try: result = subprocess.run( [sys.executable, str(TOOLS_DIR / "mutator-gate.py"), "--runtime", runtime], input=b'{"tool_name":"Read"}\n', capture_output=True, env=environment, check=False, timeout=5, ) except subprocess.TimeoutExpired as exc: server.close() thread.join(timeout=2) self.fail( f"{runtime} gate hung on wire reply {wire_reply!r}: {exc}" ) thread.join(timeout=2) self.assertFalse(thread.is_alive()) self.assertEqual(result.returncode, 2) self.assertIn(b"GATE_UNAVAILABLE", result.stderr) def test_launcher_entrypoint_returns_usage_without_a_command(self) -> None: with patch.object( sys, "argv", [str(TOOLS_DIR / "launch-runtime.py"), "--runtime", "claude"], ), redirect_stderr(io.StringIO()), self.assertRaises(SystemExit) as raised: runpy.run_path(str(TOOLS_DIR / "launch-runtime.py"), run_name="__main__") self.assertEqual(raised.exception.code, 64) def test_revoker_entrypoint_denies_when_identity_environment_is_absent(self) -> None: with patch.object( sys, "argv", [ str(TOOLS_DIR / "revoke-lease.py"), "--runtime", "claude", "--reason", "pre-compact", ], ), patch.dict(os.environ, {}, clear=True), redirect_stderr( io.StringIO() ), self.assertRaises(SystemExit) as raised: runpy.run_path(str(TOOLS_DIR / "revoke-lease.py"), run_name="__main__") self.assertEqual(raised.exception.code, 2) def test_gate_entrypoint_denies_when_identity_environment_is_absent(self) -> None: class Stdin: buffer = io.BytesIO(b'{"tool_name":"Bash"}') with patch.object( sys, "argv", [str(TOOLS_DIR / "mutator-gate.py"), "--runtime", "claude"], ), patch.object(sys, "stdin", Stdin()), patch.dict( os.environ, {}, clear=True ), redirect_stderr(io.StringIO()), self.assertRaises(SystemExit) as raised: runpy.run_path(str(TOOLS_DIR / "mutator-gate.py"), run_name="__main__") self.assertEqual(raised.exception.code, 2) class MutatorGateTest(unittest.TestCase): @staticmethod def environment() -> dict[str, str]: return { "MOSAIC_LEASE_BROKER_SOCKET": "/run/test/broker.sock", "MOSAIC_LEASE_SESSION_ID": "d" * 64, "MOSAIC_RUNTIME_GENERATION": "2", } def run_main(self, *, tool: object = "Bash", reply: dict[str, object] | None = None): calls: list[tuple[Path, dict[str, object]]] = [] def request(path: Path, payload: dict[str, object]) -> dict[str, object]: calls.append((path, payload)) return reply if reply is not None else {"ok": True, "decision": "allow"} stderr = io.StringIO() with redirect_stderr(stderr): result = GATE.main( ["--runtime", "claude"], environ=self.environment(), stream=io.BytesIO(json.dumps({"tool_name": tool}).encode()), request=request, ) return result, stderr.getvalue(), calls def test_generation_file_is_the_effective_generation_authority(self) -> None: calls: list[dict[str, object]] = [] result = GATE.main( ["--runtime", "pi"], environ=self.environment(), stream=io.BytesIO(b'{"tool_name":"bash"}'), request=lambda _path, payload: calls.append(payload) or {"ok": True, "decision": "allow"}, resolve_generation=lambda _environment: 9, ) self.assertEqual(result, 0) self.assertEqual(calls[0]["runtime_generation"], 9) def test_allow_and_denial_decisions(self) -> None: allowed, allowed_stderr, calls = self.run_main() self.assertEqual(allowed, 0) self.assertEqual(allowed_stderr, "") self.assertEqual(calls[0][0], Path("/run/test/broker.sock")) self.assertEqual( calls[0][1], { "action": "authorize_tool", "session_id": "d" * 64, "runtime_generation": 2, "runtime": "claude", "tool_name": "Bash", }, ) denied, denied_stderr, _ = self.run_main(reply={"ok": False, "code": "LEASE_EXPIRED"}) self.assertEqual(denied, 2) self.assertIn("LEASE_EXPIRED", denied_stderr) defaulted, defaulted_stderr, _ = self.run_main(reply={"ok": False, "code": 4}) self.assertEqual(defaulted, 2) self.assertIn("MUTATOR_UNVERIFIED", defaulted_stderr) def test_input_validation_fails_closed(self) -> None: payloads = [ b"x" * (GATE.MAX_FRAME + 1), b"[]", b"{}", json.dumps({"tool_name": ""}).encode(), json.dumps({"tool_name": 4}).encode(), json.dumps({"tool_name": "x" * 257}).encode(), b"not-json", ] for payload in payloads: with self.subTest(size=len(payload)), redirect_stderr(io.StringIO()): self.assertEqual( GATE.main( ["--runtime", "pi"], environ=self.environment(), stream=io.BytesIO(payload), request=lambda *_args: self.fail("invalid input reached broker"), ), 2, ) def test_environment_generation_and_request_failures_deny(self) -> None: environments = [ {}, {**self.environment(), "MOSAIC_RUNTIME_GENERATION": "bad"}, {**self.environment(), "MOSAIC_RUNTIME_GENERATION": "-1"}, ] for environment in environments: with self.subTest(environment=environment), redirect_stderr(io.StringIO()): self.assertEqual( GATE.main( ["--runtime", "claude"], environ=environment, stream=io.BytesIO(b'{"tool_name":"Read"}'), request=lambda *_args: {}, ), 2, ) failures = [ValueError("bad"), OSError("down"), json.JSONDecodeError("bad", "x", 0)] for failure in failures: with self.subTest(failure=type(failure).__name__), redirect_stderr(io.StringIO()): def request(*_args, error=failure): raise error self.assertEqual( GATE.main( ["--runtime", "claude"], environ=self.environment(), stream=io.BytesIO(b'{"tool_name":"Read"}'), request=request, ), 2, ) def test_broker_request_framing_payload_and_shape_validation(self) -> None: with self.assertRaises(ValueError): GATE.broker_request(Path("/broker"), {"session_id": "x" * GATE.MAX_FRAME}) replies = [ (b'{"ok":true,"decision":"allow"}\n', {"ok": True, "decision": "allow"}), (b'{"ok":true}', ValueError), (b'[]\n', ValueError), (b"x" * (GATE.MAX_FRAME + 1), ValueError), ] for wire_reply, expected in replies: with self.subTest(size=len(wire_reply)): fake = FakeSocket(wire_reply) with patch.object(GATE.socket, "socket", return_value=fake): if isinstance(expected, type) and issubclass(expected, Exception): with self.assertRaises(expected): GATE.broker_request(Path("/broker"), {"action": "authorize_tool"}) else: self.assertEqual( GATE.broker_request(Path("/broker"), {"action": "authorize_tool"}), expected, ) self.assertEqual(fake.timeout, GATE.BROKER_TIMEOUT_SECONDS) self.assertEqual(fake.connected, "/broker") self.assertEqual(fake.shutdown_how, socket.SHUT_WR) class LeaseRevocationTest(unittest.TestCase): def load_modules(self): generation_path = TOOLS_DIR / "lease_generation.py" revoker_path = TOOLS_DIR / "revoke-lease.py" self.assertTrue(generation_path.is_file(), "lease_generation.py must be shipped") self.assertTrue(revoker_path.is_file(), "revoke-lease.py must be shipped") return ( load_tool("lease_generation_test", "lease_generation.py"), load_tool("lease_revoker_test", "revoke-lease.py"), ) def test_generation_parser_and_descriptor_security_rejections(self) -> None: generation, _ = self.load_modules() for value in (None, "", "-1", "no", "é", str(generation.MAX_GENERATION + 1)): with self.subTest(value=value), self.assertRaises(ValueError): generation.parse_generation(value) for value in (-1, generation.MAX_GENERATION + 1): with self.subTest(initialize=value), tempfile.TemporaryDirectory() as directory, self.assertRaises(ValueError): generation.initialize_runtime_generation(Path(directory) / "state", value) with patch.object( generation.os, "fstat", return_value=SimpleNamespace(st_mode=stat.S_IFDIR | 0o700, st_uid=os.geteuid(), st_size=0), ), self.assertRaises(ValueError): generation._validate_descriptor(4) with patch.object( generation.os, "fstat", return_value=SimpleNamespace(st_mode=stat.S_IFREG | 0o600, st_uid=os.geteuid() + 1, st_size=0), ), self.assertRaises(ValueError): generation._validate_descriptor(4) with patch.object( generation.os, "fstat", return_value=SimpleNamespace(st_mode=stat.S_IFREG | 0o644, st_uid=os.geteuid(), st_size=0), ), self.assertRaises(ValueError): generation._validate_descriptor(4) with patch.object( generation.os, "fstat", return_value=SimpleNamespace( st_mode=stat.S_IFREG | 0o600, st_uid=os.geteuid(), st_size=generation.MAX_GENERATION_BYTES + 1, ), ), self.assertRaises(ValueError): generation._validate_descriptor(4) def test_generation_file_is_private_monotonic_and_rejects_unsafe_state(self) -> None: generation, _ = self.load_modules() with tempfile.TemporaryDirectory() as directory: path = Path(directory) / "runtime.generation" generation.initialize_runtime_generation(path, 4) self.assertEqual(generation.read_runtime_generation({ "MOSAIC_RUNTIME_GENERATION": "1", "MOSAIC_LEASE_GENERATION_FILE": str(path), }), 4) self.assertEqual(generation.bump_runtime_generation({ "MOSAIC_RUNTIME_GENERATION": "1", "MOSAIC_LEASE_GENERATION_FILE": str(path), }), 5) self.assertEqual(path.read_text(), "5\n") self.assertEqual(stat.S_IMODE(path.stat().st_mode), 0o600) path.write_text("bad\n") with self.assertRaises(ValueError): generation.read_runtime_generation({ "MOSAIC_RUNTIME_GENERATION": "1", "MOSAIC_LEASE_GENERATION_FILE": str(path), }) path.write_bytes(b"\xff\n") with self.assertRaises(ValueError): generation.read_runtime_generation({ "MOSAIC_LEASE_GENERATION_FILE": str(path), }) path.write_text(f"{generation.MAX_GENERATION}\n") with self.assertRaises(ValueError): generation.bump_runtime_generation({ "MOSAIC_LEASE_GENERATION_FILE": str(path), }) with self.assertRaises(ValueError): generation.bump_runtime_generation({"MOSAIC_RUNTIME_GENERATION": "1"}) def test_generation_write_must_make_progress(self) -> None: generation, _ = self.load_modules() with tempfile.TemporaryDirectory() as directory, patch.object( generation.os, "write", return_value=0 ), self.assertRaises(OSError): generation.initialize_runtime_generation(Path(directory) / "state", 1) def test_revoker_reuses_broker_revoke_and_optional_generation_bump(self) -> None: _, revoker = self.load_modules() with tempfile.TemporaryDirectory() as directory: generation_file = Path(directory) / "runtime.generation" generation_file.write_text("2\n") generation_file.chmod(0o600) environment = { "MOSAIC_LEASE_BROKER_SOCKET": "/broker", "MOSAIC_LEASE_SESSION_ID": "a" * 64, "MOSAIC_RUNTIME_GENERATION": "2", "MOSAIC_LEASE_GENERATION_FILE": str(generation_file), } calls: list[tuple[Path, dict[str, object]]] = [] result = revoker.main( ["--runtime", "pi", "--reason", "session-start-resume", "--bump-generation"], environ=environment, request=lambda path, payload: calls.append((path, payload)) or {"ok": True, "state": "UNVERIFIED"}, ) self.assertEqual(result, 0) self.assertEqual(generation_file.read_text(), "3\n") self.assertEqual(calls[0][0], Path("/broker")) self.assertEqual(calls[0][1], { "action": "revoke_lease", "session_id": "a" * 64, "runtime_generation": 3, "reason": "session-start-resume", "runtime": "pi", }) calls.clear() self.assertEqual( revoker.main( ["--runtime", "pi", "--reason", "pi-context-after-compact"], environ=environment, request=lambda path, payload: calls.append((path, payload)) or {"ok": True, "state": "UNVERIFIED"}, ), 0, ) self.assertEqual(calls[0][1]["runtime_generation"], 3) def test_revoker_broker_framing_and_shape_validation(self) -> None: _, revoker = self.load_modules() with self.assertRaises(ValueError): revoker.broker_request(Path("/broker"), {"reason": "x" * revoker.MAX_FRAME}) replies = [ (b'{"ok":true,"state":"UNVERIFIED"}\n', {"ok": True, "state": "UNVERIFIED"}), (b'{"ok":true}', ValueError), (b'[]\n', ValueError), (b'x' * (revoker.MAX_FRAME + 1), ValueError), ] for wire_reply, expected in replies: with self.subTest(size=len(wire_reply)): fake = FakeSocket(wire_reply) with patch.object(revoker.socket, "socket", return_value=fake): if isinstance(expected, type) and issubclass(expected, Exception): with self.assertRaises(expected): revoker.broker_request(Path("/broker"), {"action": "revoke_lease"}) else: self.assertEqual( revoker.broker_request(Path("/broker"), {"action": "revoke_lease"}), expected, ) self.assertEqual(fake.timeout, revoker.BROKER_TIMEOUT_SECONDS) self.assertEqual(fake.connected, "/broker") self.assertEqual(fake.shutdown_how, socket.SHUT_WR) def test_failed_observer_revocation_advances_the_local_generation_fence(self) -> None: _, revoker = self.load_modules() with tempfile.TemporaryDirectory() as directory: generation_file = Path(directory) / "runtime.generation" generation_file.write_text("8\n") generation_file.chmod(0o600) environment = { "MOSAIC_LEASE_BROKER_SOCKET": "/broker", "MOSAIC_LEASE_SESSION_ID": "a" * 64, "MOSAIC_RUNTIME_GENERATION": "8", "MOSAIC_LEASE_GENERATION_FILE": str(generation_file), } with redirect_stderr(io.StringIO()): result = revoker.main( ["--runtime", "claude", "--reason", "session-start-compact"], environ=environment, request=lambda *_args: (_ for _ in ()).throw(OSError("down")), ) self.assertEqual(result, 2) self.assertEqual(generation_file.read_text(), "9\n") def test_revoker_fails_closed_on_identity_reply_and_transport_errors(self) -> None: _, revoker = self.load_modules() good = { "MOSAIC_LEASE_BROKER_SOCKET": "/broker", "MOSAIC_LEASE_SESSION_ID": "a" * 64, "MOSAIC_RUNTIME_GENERATION": "1", } malformed_session = {**good, "MOSAIC_LEASE_SESSION_ID": "not-a-session"} cases = [ ({}, lambda *_args: {"ok": True, "state": "UNVERIFIED"}), (malformed_session, lambda *_args: {"ok": True, "state": "UNVERIFIED"}), (good, lambda *_args: {"ok": False, "state": "UNVERIFIED"}), (good, lambda *_args: {"ok": True, "state": "VERIFIED"}), (good, lambda *_args: (_ for _ in ()).throw(OSError("down"))), ( good, lambda *_args: (_ for _ in ()).throw( json.JSONDecodeError("bad", "x", 0) ), ), ] for environment, request in cases: with self.subTest(environment=environment), redirect_stderr(io.StringIO()): self.assertEqual( revoker.main( ["--runtime", "claude", "--reason", "pre-compact"], environ=environment, request=request, ), 2, ) with redirect_stderr(io.StringIO()): self.assertEqual( revoker.main( ["--runtime", "claude", "--reason", "x" * 129], environ=good, request=lambda *_args: self.fail("invalid reason reached broker"), ), 2, ) with tempfile.TemporaryDirectory() as directory: generation_file = Path(directory) / "runtime.generation" generation_file.write_text("1\n") generation_file.chmod(0o600) with redirect_stderr(io.StringIO()): self.assertEqual( revoker.main( [ "--runtime", "pi", "--reason", "session-start-resume", "--bump-generation", ], environ={ **good, "MOSAIC_LEASE_GENERATION_FILE": str(generation_file), }, request=lambda *_args: (_ for _ in ()).throw(OSError("down")), ), 2, ) self.assertEqual(generation_file.read_text(), "2\n") if __name__ == "__main__": unittest.main()