from __future__ import annotations import json import stat import tempfile import unittest from dataclasses import replace from pathlib import Path from unittest.mock import patch from uvm.config import Settings from uvm.domain import STATE_SCHEMA_VERSION, PortRule, State, VmRecord, VmSpec, new_vm_id from uvm.errors import ConflictError, NotFoundError, StateError from uvm.state import StateStore def make_vm(vm_id: str = "vm-test") -> VmRecord: return VmRecord( id=vm_id, cpu=1, ram_mib=512, guest_ip="10.42.0.2", gateway="10.42.0.1", tap="uvm-test", mac="02:fc:00:00:00:01", socket="/tmp/firecracker.sock", config="/tmp/config.json", log="/tmp/firecracker.log", disk="/tmp/rootfs.ext4", username="admin", password="stored-secret", ) class StateStoreTests(unittest.TestCase): def setUp(self) -> None: self._temporary_directory = tempfile.TemporaryDirectory() self.base = Path(self._temporary_directory.name) / "uvm" self.settings = replace(Settings(), base=self.base) self.store = StateStore(self.settings) def tearDown(self) -> None: self._temporary_directory.cleanup() def test_transaction_persists_a_vm_atomically(self) -> None: self.store.initialize() with self.store.transaction() as state: state.vms["vm-test"] = make_vm() state.next_mac_index = 2 loaded = self.store.load() self.assertEqual(loaded.vms["vm-test"].disk, "/tmp/rootfs.ext4") self.assertEqual(loaded.vms["vm-test"].username, "admin") self.assertEqual(loaded.vms["vm-test"].password, "stored-secret") self.assertNotIn("stored-secret", repr(loaded.vms["vm-test"])) self.assertEqual(loaded.next_mac_index, 2) self.assertEqual(stat.S_IMODE(self.settings.state_path.stat().st_mode), 0o600) persisted = json.loads(self.settings.state_path.read_text(encoding="utf-8")) self.assertEqual(persisted["vms"]["vm-test"]["username"], "admin") self.assertEqual(persisted["vms"]["vm-test"]["password"], "stored-secret") def test_legacy_state_is_loaded_and_migrated_on_next_write(self) -> None: self.base.mkdir(parents=True) for version in (None, 1): with self.subTest(version=version): raw_vm = make_vm().to_dict() for field in ( "disk", "username", "password", "updated_at", "name", "ports", "disk_ready", "runtime_owned", "tap_owned", "operation", "boot_id", "spawn_pending", "network", "bridge", ): raw_vm.pop(field) legacy = {"vms": {"vm-test": raw_vm}} if version is not None: legacy["schema_version"] = version self.settings.state_path.write_text(json.dumps(legacy), encoding="utf-8") loaded = self.store.load() vm = loaded.vms["vm-test"] self.assertEqual(loaded.schema_version, 2) self.assertEqual(vm.disk, "") self.assertEqual(vm.username, "root") self.assertIsNone(vm.password) self.assertIsNone(vm.name) self.assertIsNone(vm.ports) self.assertIsNone(vm.network) self.assertIsNone(vm.bridge) self.assertFalse(vm.disk_ready) self.assertFalse(vm.runtime_owned) self.assertFalse(vm.tap_owned) self.assertFalse(vm.spawn_pending) self.assertEqual(vm.updated_at, vm.created_at) self.assertEqual(loaded.next_mac_index, 2) self.assertEqual( json.loads(self.settings.state_path.read_text(encoding="utf-8")), legacy ) with self.store.transaction() as state: state.vms["vm-test"].status = "stopped" persisted = json.loads(self.settings.state_path.read_text(encoding="utf-8")) self.assertEqual(persisted["schema_version"], 2) self.assertEqual(persisted["next_mac_index"], 2) self.assertEqual(persisted["vms"]["vm-test"]["status"], "stopped") self.assertIsNone(persisted["vms"]["vm-test"]["ports"]) def test_round_trip_preserves_policy_ownership_and_recovery_fields(self) -> None: for ports in (None, (), (PortRule(22), PortRule(8080, "http"))): with self.subTest(ports=ports): vm = replace( make_vm(), name="app-server", ports=ports, disk_ready=True, runtime_owned=True, tap_owned=True, operation="configuring", boot_id="host-boot-id", spawn_pending=True, pid=123, process_start_time="456789", network="10.42.0.0/24", bridge="uvm0", last_error="interrupted boot", ) with self.store.transaction() as state: state.vms[vm.id] = vm loaded = self.store.load() self.assertEqual(loaded.schema_version, 2) self.assertEqual(loaded.vms[vm.id], vm) raw = json.loads(self.settings.state_path.read_text(encoding="utf-8")) self.assertEqual(raw["schema_version"], 2) self.assertEqual( raw["vms"][vm.id]["ports"], None if ports is None else [port.to_dict() for port in ports], ) for field in ("disk_ready", "runtime_owned", "tap_owned", "spawn_pending"): self.assertIs(raw["vms"][vm.id][field], True) def test_lookup_requires_full_case_sensitive_ids_names_or_ips(self) -> None: vm = replace(make_vm("vm-0123456789abcdef0123456789abcdef"), name="app-server") with self.store.transaction() as state: state.vms[vm.id] = vm for identifier in (vm.id, vm.name, vm.guest_ip): with self.subTest(identifier=identifier): self.assertEqual(self.store.find_by_id_or_ip(identifier), vm) self.assertEqual(self.store.get(vm.id), vm) for identifier in ( "", vm.id[:15], vm.id.upper(), "app", "APP-SERVER", " app-server ", "10.42.0", ): with self.subTest(identifier=identifier), self.assertRaises(NotFoundError): self.store.find_by_id_or_ip(identifier) for identifier in (vm.name, vm.guest_ip): with self.subTest(get_identifier=identifier), self.assertRaises(NotFoundError): self.store.get(identifier) def test_load_rejects_a_record_whose_id_does_not_match_its_key(self) -> None: self.base.mkdir(parents=True) for version in (0, 1, 2): with self.subTest(version=version): raw = { "schema_version": version, "vms": {"vm-wrong-key": make_vm("vm-record-id").to_dict()}, } self.settings.state_path.write_text(json.dumps(raw), encoding="utf-8") with self.assertRaisesRegex(StateError, "ID.*registry key"): self.store.load() def test_initialize_protects_registry_credentials(self) -> None: self.base.mkdir(parents=True) self.settings.state_path.write_text('{"vms": {}}', encoding="utf-8") self.settings.state_path.chmod(0o644) self.store.initialize() self.assertEqual(stat.S_IMODE(self.settings.state_path.stat().st_mode), 0o600) def test_load_protects_registry_credentials_without_initialization(self) -> None: self.base.mkdir(parents=True) self.settings.state_path.write_text('{"vms": {}}', encoding="utf-8") self.settings.state_path.chmod(0o644) self.store.load() self.assertEqual(stat.S_IMODE(self.settings.state_path.stat().st_mode), 0o600) def test_load_explains_when_registry_permissions_require_root(self) -> None: self.base.mkdir(parents=True) self.settings.state_path.write_text('{"vms": {}}', encoding="utf-8") self.settings.state_path.chmod(0o644) with ( patch("uvm.state.os.chmod", side_effect=PermissionError), self.assertRaisesRegex(StateError, "Run this command with sudo"), ): self.store.load() class DomainStateTests(unittest.TestCase): def test_new_specs_default_to_ssh_without_overriding_an_explicit_empty_policy(self) -> None: self.assertEqual(VmSpec(cpu=1, ram_mib=512).ports, (PortRule(22),)) self.assertEqual(VmSpec(cpu=1, ram_mib=512, ports=()).ports, ()) self.assertIsNone(make_vm().ports) def test_port_rule_serialization_normalizes_labels_and_defaults_to_tcp(self) -> None: for port, protocol in ((1, "tcp"), (8080, "hTtP"), (65535, "https")): with self.subTest(port=port, protocol=protocol): rule = PortRule.from_dict({"port": port, "protocol": protocol}) self.assertEqual(rule.to_dict(), {"port": port, "protocol": protocol.upper()}) self.assertEqual(PortRule.from_dict(rule.to_dict()), rule) self.assertEqual(PortRule.from_dict({"port": 22}), PortRule(22, "TCP")) def test_port_rule_deserialization_rejects_invalid_types_fields_and_values(self) -> None: invalid = ( None, True, 22, "22", [], (), {}, {"protocol": "TCP"}, {"port": True}, {"port": False}, {"port": None}, {"port": "22"}, {"port": 22.0}, {"port": 0}, {"port": 65536}, {"port": 22, "protocol": None}, {"port": 22, "protocol": True}, {"port": 22, "protocol": "UDP"}, {"port": 22, "protocol": " TCP"}, {"port": 22, "type": "TCP"}, {"port": 22, "protocol": "TCP", "extra": 1}, ) for value in invalid: with self.subTest(value=value), self.assertRaises(StateError): PortRule.from_dict(value) def test_persisted_ports_must_be_a_json_list_or_legacy_null(self) -> None: for ports in (True, False, 22, "22", {}, (), (PortRule(22),)): with self.subTest(ports=ports), self.assertRaisesRegex(StateError, "ports"): VmRecord.from_dict({**make_vm().to_dict(), "ports": ports}) for ports in ([22], [None], ["22:TCP"], [{"port": True}]): with self.subTest(ports=ports), self.assertRaises(StateError): VmRecord.from_dict({**make_vm().to_dict(), "ports": ports}) def test_persisted_duplicate_and_conflicting_tcp_rules_are_rejected(self) -> None: for labels in (("TCP", "TCP"), ("tcp", "TCP"), ("TCP", "HTTP"), ("HTTP", "HTTPS")): with self.subTest(labels=labels): raw = make_vm().to_dict() raw["ports"] = [{"port": 80, "protocol": label} for label in labels] with self.assertRaisesRegex(StateError, "duplicate TCP port"): VmRecord.from_dict(raw) def test_ownership_and_spawn_pending_fields_require_real_booleans(self) -> None: for field in ("disk_ready", "runtime_owned", "tap_owned", "spawn_pending"): for value in (None, 0, 1, 0.0, 1.0, "false", "true", [], {}): with self.subTest(field=field, value=value): raw = {**make_vm().to_dict(), field: value} with self.assertRaisesRegex(StateError, f"{field} must be a boolean"): VmRecord.from_dict(raw) def test_explicit_recovery_flags_override_legacy_inference(self) -> None: fields = ("disk_ready", "runtime_owned", "tap_owned", "spawn_pending") for flag in (False, True): for operation in (None, "spawning", "configuring"): with self.subTest(flag=flag, operation=operation): raw = replace(make_vm(), status="running", pid=123, operation=operation).to_dict() raw.update(dict.fromkeys(fields, flag)) vm = VmRecord.from_dict(raw) for field in fields: self.assertIs(getattr(vm, field), flag) self.assertIs(vm.to_dict()[field], flag) def test_legacy_ownership_is_inferred_from_disk_status_and_pid(self) -> None: cases = ( ("running", "/tmp/rootfs.ext4", 123, (True, True, True)), ("stopped", "/tmp/rootfs.ext4", None, (True, True, False)), ("error", "/tmp/rootfs.ext4", 123, (False, True, True)), ("starting", "/tmp/rootfs.ext4", None, (False, True, False)), ("running", "", 123, (False, False, True)), ("unknown", "", None, (False, False, False)), ) for status, disk, pid, expected in cases: with self.subTest(status=status, disk=disk, pid=pid): raw = replace(make_vm(), status=status, disk=disk, pid=pid).to_dict() for field in ("disk_ready", "runtime_owned", "tap_owned", "ports"): raw.pop(field) vm = VmRecord.from_dict(raw) self.assertEqual((vm.disk_ready, vm.runtime_owned, vm.tap_owned), expected) self.assertIsNone(vm.ports) def test_missing_spawn_pending_recovers_only_a_legacy_spawning_operation(self) -> None: for operation in (None, "spawning", "configuring", "stopping", "policy"): with self.subTest(operation=operation): raw = replace(make_vm(), operation=operation).to_dict() raw.pop("spawn_pending") vm = VmRecord.from_dict(raw) self.assertIs(vm.spawn_pending, operation == "spawning") self.assertEqual(VmRecord.from_dict(vm.to_dict()), vm) def test_new_state_uses_schema_two_and_rejects_future_versions(self) -> None: self.assertEqual(State().schema_version, 2) with self.assertRaisesRegex(StateError, "newer than"): State.from_dict({"schema_version": STATE_SCHEMA_VERSION + 1, "vms": {}}) def test_generated_ids_have_the_reserved_uuid_format(self) -> None: first, second = new_vm_id(), new_vm_id() self.assertNotEqual(first, second) for vm_id in (first, second): self.assertRegex(vm_id, r"\Avm-[0-9a-f]{32}\Z") def test_exact_generated_id_wins_over_a_legacy_name_collision(self) -> None: vm = make_vm("vm-" + "a" * 32) other = replace(make_vm("vm-" + "b" * 32), name=vm.id, guest_ip="10.42.0.3") state = State(vms={other.id: other, vm.id: vm}) self.assertIs(StateStore.resolve(state, vm.id), vm) del state.vms[vm.id] with self.assertRaises(NotFoundError): StateStore.resolve(state, vm.id) def test_ambiguous_names_ips_and_legacy_ids_require_an_exact_id(self) -> None: cases = ( ("vm-" + "a" * 32, "app", "app", "10.42.0.3", "app"), ("vm-" + "a" * 32, "app", "other", "10.42.0.2", "10.42.0.2"), ("vm-" + "a" * 32, "app", "10.42.0.2", "10.42.0.3", "10.42.0.2"), ("vm-legacy", "app", "vm-legacy", "10.42.0.3", "vm-legacy"), ) for vm_id, name, other_name, other_ip, identifier in cases: with self.subTest(identifier=identifier, other_name=other_name): vm = replace(make_vm(vm_id), name=name) other = replace(make_vm("vm-" + "b" * 32), name=other_name, guest_ip=other_ip) for records in ((vm, other), (other, vm)): state = State(vms={record.id: record for record in records}) with self.assertRaisesRegex(ConflictError, "ambiguous.*exact VM ID"): StateStore.resolve(state, identifier) self.assertIs(StateStore.resolve(state, other.id), other) def test_multiple_matching_fields_on_one_vm_are_not_ambiguous(self) -> None: vm = replace(make_vm("vm-legacy"), name="vm-legacy") self.assertIs(StateStore.resolve(State(vms={vm.id: vm}), "vm-legacy"), vm)