nest/tests/test_network.py

from __future__ import annotations

import unittest
from ipaddress import IPv4Address, IPv4Network

from uvm.config import Settings
from uvm.domain import VmRecord
from uvm.errors import UvmError, ValidationError
from uvm.network import NetworkManager
from uvm.system import CommandResult


class FakeRunner:
    def __init__(self) -> None:
        self.calls: list[tuple[str, ...]] = []
        self.existing_links: set[str] = set()
        self.link_details: dict[str, str] = {}

    def run(self, command, *, check=True, capture=False, timeout=None) -> CommandResult:
        del check, capture, timeout
        args = tuple(str(part) for part in command)
        self.calls.append(args)
        if args[:3] == ("ip", "link", "show") and args[3] in self.existing_links:
            return CommandResult(args=args, returncode=0)
        if args[:6] == ("ip", "-j", "-d", "link", "show", "dev"):
            return CommandResult(args=args, returncode=0, stdout=self.link_details[args[6]])
        if args[:4] == ("ip", "link", "show", "uvm0"):
            return CommandResult(args=args, returncode=1)
        if args == ("ip", "route", "show", "default"):
            return CommandResult(args=args, returncode=0, stdout="default via 192.0.2.1 dev eth0\n")
        if len(args) > 3 and args[0] == "iptables" and "-C" in args:
            return CommandResult(args=args, returncode=1)
        return CommandResult(args=args, returncode=0)


class NetworkManagerTests(unittest.TestCase):
    def setUp(self) -> None:
        self.runner = FakeRunner()
        self.network = NetworkManager(Settings(), runner=self.runner)  # type: ignore[arg-type]

    def test_allocates_first_available_guest_address(self) -> None:
        self.assertEqual(self.network.allocate_ip([], None), IPv4Address("10.42.0.2"))

    def test_rejects_gateway_and_used_requested_addresses(self) -> None:
        vm = 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",
        )
        with self.assertRaises(ValidationError):
            self.network.allocate_ip([vm], IPv4Address("10.42.0.1"))
        with self.assertRaises(ValidationError):
            self.network.allocate_ip([vm], IPv4Address("10.42.0.2"))
        with self.assertRaises(ValidationError):
            self.network.allocate_ip([], IPv4Address("10.42.0.0"))
        with self.assertRaises(ValidationError):
            self.network.allocate_ip([], IPv4Address("10.42.0.255"))

    def test_mac_and_tap_names_are_bounded_and_deterministic(self) -> None:
        self.assertEqual(NetworkManager.mac_for(1), "02:fc:00:00:00:01")
        tap = NetworkManager.tap_name("vm-0123456789abcdef")
        self.assertLessEqual(len(tap), 15)
        self.assertTrue(tap.startswith("uvm-"))

    def test_bridge_setup_adds_a_missing_bridge_and_one_nat_rule(self) -> None:
        self.network.ensure_bridge()

        self.assertIn(("ip", "link", "add", "uvm0", "type", "bridge"), self.runner.calls)
        self.assertIn(
            ("ip", "addr", "replace", "10.42.0.1/24", "dev", "uvm0"),
            self.runner.calls,
        )
        self.assertIn(
            (
                "iptables",
                "-A",
                "FORWARD",
                "-i",
                "uvm0",
                "-o",
                "eth0",
                "-s",
                "10.42.0.0/24",
                "-j",
                "ACCEPT",
            ),
            self.runner.calls,
        )
        checks = [
            call
            for call in self.runner.calls
            if call[:5] == ("iptables", "-t", "nat", "-C", "POSTROUTING")
        ]
        self.assertEqual(len(checks), 1)
        self.assertIn(
            (
                "iptables",
                "-t",
                "nat",
                "-A",
                "POSTROUTING",
                "-s",
                "10.42.0.0/24",
                "-o",
                "eth0",
                "-j",
                "MASQUERADE",
            ),
            self.runner.calls,
        )

    def test_bridge_and_guest_network_support_non_default_prefixes(self) -> None:
        settings = Settings(
            network=IPv4Network("10.50.0.0/16"),
            gateway=IPv4Address("10.50.0.1"),
        )
        runner = FakeRunner()
        NetworkManager(settings, runner=runner).ensure_bridge()  # type: ignore[arg-type]

        self.assertIn(
            ("ip", "addr", "replace", "10.50.0.1/16", "dev", "uvm0"),
            runner.calls,
        )

    def test_refuses_to_adopt_an_existing_tap_device(self) -> None:
        self.runner.existing_links.add("uvm-existing")

        with self.assertRaises(UvmError):
            self.network.create_tap("uvm-existing")

        self.assertNotIn(
            ("ip", "tuntap", "add", "dev", "uvm-existing", "mode", "tap"),
            self.runner.calls,
        )

    def test_refuses_to_reconfigure_an_existing_non_bridge_interface(self) -> None:
        self.runner.existing_links.add("uvm0")
        self.runner.link_details["uvm0"] = '[{"linkinfo": {"info_kind": "dummy"}}]'

        with self.assertRaises(UvmError):
            self.network.ensure_bridge()

        self.assertNotIn(
            ("ip", "addr", "replace", "10.42.0.1/24", "dev", "uvm0"),
            self.runner.calls,
        )