#!/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 sys import unittest from contextlib import redirect_stderr from pathlib import Path from unittest.mock import patch TOOLS_DIR = Path(__file__).parents[2] / "framework/tools/lease-broker" 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) 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, ) 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["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), ) 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), ) 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_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")), ), 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_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_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_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) if __name__ == "__main__": unittest.main()