100 lines
3.7 KiB
Python
100 lines
3.7 KiB
Python
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)
|
|
|
|
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)
|