Files
2026-09-04 20:56:18 +00:00

161 lines
5.1 KiB
Python

from __future__ import annotations
import argparse
import io
import unittest
from contextlib import redirect_stdout
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch
from uvm.cli import _run_command
from uvm.domain import VmRecord
from uvm.images import GuestAssets
def make_vm() -> VmRecord:
return VmRecord(
id="vm-test",
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",
username="admin",
password="stored-secret",
)
class CliSshTests(unittest.TestCase):
def setUp(self) -> None:
self.vm = make_vm()
self.application = SimpleNamespace(
lifecycle=SimpleNamespace(find_for_ssh=lambda _identifier: self.vm)
)
def test_ssh_uses_a_stable_vm_host_key_alias_by_default(self) -> None:
args = argparse.Namespace(
command="ssh",
vm=self.vm.id,
key=None,
user="root",
insecure_host_key=False,
)
with patch("uvm.cli.os.execvp") as execvp:
_run_command(self.application, args)
command = execvp.call_args.args[1]
self.assertIn("HostKeyAlias=uvm-vm-test", command)
self.assertIn("StrictHostKeyChecking=accept-new", command)
self.assertNotIn("StrictHostKeyChecking=no", command)
def test_ssh_only_disables_verification_when_explicitly_requested(self) -> None:
args = argparse.Namespace(
command="ssh",
vm=self.vm.id,
key=None,
user="root",
insecure_host_key=True,
)
with patch("uvm.cli.os.execvp") as execvp:
_run_command(self.application, args)
command = execvp.call_args.args[1]
self.assertIn("StrictHostKeyChecking=no", command)
self.assertNotIn("HostKeyAlias=uvm-vm-test", command)
def test_ssh_uses_the_username_stored_for_the_vm(self) -> None:
args = argparse.Namespace(
command="ssh",
vm=self.vm.id,
key=None,
user=None,
insecure_host_key=False,
)
with patch("uvm.cli.os.execvp") as execvp:
_run_command(self.application, args)
self.assertEqual(execvp.call_args.args[1][-1], "admin@10.42.0.2")
class CliCreateTests(unittest.TestCase):
def test_create_passes_credentials_without_printing_the_password(self) -> None:
captured = None
def create(spec):
nonlocal captured
captured = spec
vm = make_vm()
vm.username = spec.username
vm.password = spec.password
return vm
application = SimpleNamespace(lifecycle=SimpleNamespace(create=create))
args = argparse.Namespace(
command="create",
cpu="1",
ram="512",
host_ip=None,
username="admin",
password="custom-secret",
)
output = io.StringIO()
with redirect_stdout(output):
result = _run_command(application, args)
self.assertEqual(result, 0)
self.assertEqual(captured.username, "admin")
self.assertEqual(captured.password, "custom-secret")
self.assertNotIn("custom-secret", output.getvalue())
class CliInstallTests(unittest.TestCase):
def test_install_next_step_does_not_assume_an_installed_console_command(self) -> None:
state_store = SimpleNamespace(initialize=lambda: None)
application = SimpleNamespace(
settings=SimpleNamespace(allow_unverified_downloads=False),
installer=SimpleNamespace(
install=lambda **_kwargs: (
Path("/tmp/firecracker"),
GuestAssets(kernel=Path("/tmp/vmlinux"), rootfs=Path("/tmp/rootfs")),
)
),
state_store=state_store,
)
args = argparse.Namespace(command="install", force=False)
output = io.StringIO()
with redirect_stdout(output):
result = _run_command(application, args)
self.assertEqual(result, 0)
self.assertIn("Run your UVM command with: create --cpu 1 --ram 512", output.getvalue())
self.assertNotIn("sudo uvm create", output.getvalue())
def test_unverified_install_reminds_the_operator_to_keep_the_opt_out(self) -> None:
application = SimpleNamespace(
settings=SimpleNamespace(allow_unverified_downloads=True),
installer=SimpleNamespace(
install=lambda **_kwargs: (
Path("/tmp/firecracker"),
GuestAssets(kernel=Path("/tmp/vmlinux"), rootfs=Path("/tmp/rootfs")),
)
),
state_store=SimpleNamespace(initialize=lambda: None),
)
args = argparse.Namespace(command="install", force=False)
output = io.StringIO()
with redirect_stdout(output):
_run_command(application, args)
self.assertIn("UVM_ALLOW_UNVERIFIED_DOWNLOADS=1", output.getvalue())