__init__

tests/test_state.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+import json
+import stat
+import tempfile
+import unittest
+from dataclasses import replace
+from pathlib import Path
+
+from uvm.config import Settings
+from uvm.domain import VmRecord
+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)
+        legacy = {"vms": {"vm-test": make_vm().to_dict()}}
+        legacy["vms"]["vm-test"].pop("disk")
+        legacy["vms"]["vm-test"].pop("username")
+        legacy["vms"]["vm-test"].pop("password")
+        legacy["vms"]["vm-test"].pop("updated_at")
+        self.settings.state_path.write_text(json.dumps(legacy), encoding="utf-8")
+
+        loaded = self.store.load()
+        self.assertEqual(loaded.vms["vm-test"].disk, "")
+        self.assertEqual(loaded.vms["vm-test"].username, "root")
+        self.assertIsNone(loaded.vms["vm-test"].password)
+        self.assertEqual(loaded.next_mac_index, 2)
+
+        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"], 1)
+        self.assertEqual(persisted["next_mac_index"], 2)
+        self.assertEqual(persisted["vms"]["vm-test"]["status"], "stopped")
+
+    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)