186 lines
6.5 KiB
Python
186 lines
6.5 KiB
Python
from __future__ import annotations
|
|
|
|
import tempfile
|
|
import unittest
|
|
from dataclasses import replace
|
|
from ipaddress import IPv4Address
|
|
from pathlib import Path
|
|
from unittest.mock import Mock, patch
|
|
|
|
from uvm.config import Settings
|
|
from uvm.domain import VmSpec
|
|
from uvm.errors import FirecrackerError, UvmError
|
|
from uvm.firecracker.process import ProcessInfo
|
|
from uvm.images import ImageStore
|
|
from uvm.lifecycle import LifecycleService
|
|
from uvm.state import StateStore
|
|
|
|
|
|
class FakeNetwork:
|
|
def __init__(self) -> None:
|
|
self.created_taps: list[str] = []
|
|
self.deleted_taps: list[str] = []
|
|
|
|
def allocate_ip(self, _vms, requested: IPv4Address | None) -> IPv4Address:
|
|
return requested or IPv4Address("10.42.0.2")
|
|
|
|
@staticmethod
|
|
def mac_for(index: int) -> str:
|
|
return f"02:fc:00:00:00:{index:02x}"
|
|
|
|
@staticmethod
|
|
def tap_name(vm_id: str) -> str:
|
|
return f"uvm-{vm_id[-8:]}"[:15]
|
|
|
|
def ensure_bridge(self) -> None:
|
|
return None
|
|
|
|
def create_tap(self, name: str) -> None:
|
|
self.created_taps.append(name)
|
|
|
|
def delete_tap(self, name: str) -> None:
|
|
self.deleted_taps.append(name)
|
|
|
|
|
|
class FakeProcess:
|
|
def __init__(self) -> None:
|
|
self.started = False
|
|
self.terminated: list[int | None] = []
|
|
self.alive = True
|
|
|
|
def start(self, _vm) -> ProcessInfo:
|
|
self.started = True
|
|
return ProcessInfo(pid=12345, start_time="42")
|
|
|
|
def is_alive(self, _vm) -> bool:
|
|
return self.alive
|
|
|
|
def terminate(self, vm) -> None:
|
|
self.terminated.append(vm.pid)
|
|
self.alive = False
|
|
|
|
|
|
class FakeClient:
|
|
def __init__(self, should_fail: bool = False) -> None:
|
|
self.should_fail = should_fail
|
|
self.config = None
|
|
|
|
def configure_and_start(self, config) -> None:
|
|
self.config = config
|
|
if self.should_fail:
|
|
raise FirecrackerError("simulated API failure")
|
|
|
|
|
|
class LifecycleTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self._temporary_directory = tempfile.TemporaryDirectory()
|
|
base = Path(self._temporary_directory.name) / "uvm"
|
|
self.settings = replace(Settings(), base=base, allow_unverified_downloads=True)
|
|
self.settings.bin_dir.mkdir(parents=True)
|
|
self.settings.firecracker_binary.touch()
|
|
self.settings.images_dir.mkdir(parents=True)
|
|
self.settings.kernel_image.write_bytes(b"kernel")
|
|
self.settings.rootfs_image.write_bytes(b"rootfs")
|
|
self.store = StateStore(self.settings)
|
|
self.images = ImageStore(self.settings)
|
|
self.images.provision_credentials = Mock() # type: ignore[method-assign]
|
|
self.network = FakeNetwork()
|
|
self.process = FakeProcess()
|
|
self.client = FakeClient()
|
|
self.service = LifecycleService(
|
|
self.settings,
|
|
self.store,
|
|
self.images,
|
|
self.network, # type: ignore[arg-type]
|
|
self.process, # type: ignore[arg-type]
|
|
client_factory=lambda _socket, _timeout: self.client, # type: ignore[arg-type]
|
|
)
|
|
self._root_patch = patch("uvm.lifecycle.require_root")
|
|
self._kvm_patch = patch("uvm.lifecycle.check_kvm")
|
|
self._root_patch.start()
|
|
self._kvm_patch.start()
|
|
|
|
def tearDown(self) -> None:
|
|
self._kvm_patch.stop()
|
|
self._root_patch.stop()
|
|
self._temporary_directory.cleanup()
|
|
|
|
def test_create_uses_a_private_disk_and_persists_running_state(self) -> None:
|
|
vm = self.service.create(VmSpec(cpu=1, ram_mib=512))
|
|
|
|
self.assertEqual(vm.status, "running")
|
|
self.assertNotEqual(Path(vm.disk), self.settings.rootfs_image)
|
|
self.assertEqual(Path(vm.disk).read_bytes(), b"rootfs")
|
|
self.assertEqual(vm.username, "root")
|
|
self.assertEqual(vm.password, "root")
|
|
self.assertEqual(self.store.load().vms[vm.id].pid, 12345)
|
|
self.assertEqual(self.network.created_taps, [vm.tap])
|
|
self.assertEqual(self.client.config["drives"][0]["path_on_host"], vm.disk)
|
|
self.images.provision_credentials.assert_called_once_with(
|
|
Path(vm.disk),
|
|
"root",
|
|
"root",
|
|
)
|
|
|
|
def test_create_persists_and_provisions_custom_guest_credentials(self) -> None:
|
|
vm = self.service.create(
|
|
VmSpec(cpu=1, ram_mib=512, username="admin", password="secret-value")
|
|
)
|
|
|
|
persisted = self.store.load().vms[vm.id]
|
|
self.assertEqual(persisted.username, "admin")
|
|
self.assertEqual(persisted.password, "secret-value")
|
|
self.images.provision_credentials.assert_called_once_with(
|
|
Path(vm.disk),
|
|
"admin",
|
|
"secret-value",
|
|
)
|
|
|
|
def test_create_rolls_back_when_firecracker_configuration_fails(self) -> None:
|
|
self.client.should_fail = True
|
|
|
|
with self.assertRaises(FirecrackerError):
|
|
self.service.create(VmSpec(cpu=1, ram_mib=512))
|
|
|
|
self.assertEqual(self.store.load().vms, {})
|
|
self.assertEqual(len(self.network.deleted_taps), 1)
|
|
self.assertEqual(self.process.terminated, [12345])
|
|
self.assertEqual(list(self.settings.vms_dir.iterdir()), [])
|
|
|
|
def test_create_rolls_back_before_network_setup_when_provisioning_fails(self) -> None:
|
|
self.images.provision_credentials.side_effect = UvmError("invalid guest image")
|
|
|
|
with self.assertRaises(UvmError):
|
|
self.service.create(VmSpec(cpu=1, ram_mib=512))
|
|
|
|
self.assertEqual(self.store.load().vms, {})
|
|
self.assertEqual(self.network.created_taps, [])
|
|
self.assertFalse(self.process.started)
|
|
self.assertEqual(list(self.settings.vms_dir.iterdir()), [])
|
|
|
|
def test_create_rolls_back_when_interrupted(self) -> None:
|
|
def interrupting_create_tap(_name: str) -> None:
|
|
raise KeyboardInterrupt
|
|
|
|
self.network.create_tap = interrupting_create_tap
|
|
|
|
with self.assertRaises(KeyboardInterrupt):
|
|
self.service.create(VmSpec(cpu=1, ram_mib=512))
|
|
|
|
self.assertEqual(self.store.load().vms, {})
|
|
self.assertEqual(list(self.settings.vms_dir.iterdir()), [])
|
|
|
|
def test_stop_keeps_disk_and_destroy_releases_state(self) -> None:
|
|
vm = self.service.create(VmSpec(cpu=1, ram_mib=512))
|
|
disk = Path(vm.disk)
|
|
|
|
stopped = self.service.stop(vm.id)
|
|
self.assertEqual(stopped.status, "stopped")
|
|
self.assertTrue(disk.exists())
|
|
self.assertIsNone(stopped.pid)
|
|
|
|
destroyed = self.service.destroy(vm.id)
|
|
self.assertEqual(destroyed.id, vm.id)
|
|
self.assertFalse(self.settings.vm_dir(vm.id).exists())
|
|
self.assertEqual(self.store.load().vms, {})
|