Merge pull request #793 from evakhoni/qemu_resize

Add disk resize support for QEMU driver

Author: bennyz

packages/jumpstarter-driver-qemu/jumpstarter_driver_qemu/client.py

@@ -1,8 +1,11 @@
 from contextlib import contextmanager
 
+import click
 from jumpstarter_driver_composite.client import CompositeClient
 from jumpstarter_driver_network.adapters import FabricAdapter, NovncAdapter
 
+from jumpstarter.client.decorators import driver_click_group
+
 
 class QemuClient(CompositeClient):
     @property
@@ -17,6 +20,14 @@ def username(self) -> str:
     def password(self) -> str:
         return self.call("get_password")
 
+    def set_disk_size(self, size: str) -> None:
+        """Set the disk size for resizing before boot."""
+        self.call("set_disk_size", size)
+
+    def set_memory_size(self, size: str) -> None:
+        """Set the memory size for next boot."""
+        self.call("set_memory_size", size)
+
     @contextmanager
     def novnc(self):
         with NovncAdapter(client=self.vnc) as url:
@@ -30,3 +41,30 @@ def shell(self):
             connect_kwargs={"password": self.password},
         ) as conn:
             yield conn
+
+    def cli(self):
+        @driver_click_group(self)
+        def base():
+            """QEMU virtual machine operations"""
+            pass
+
+        @base.group()
+        def resize():
+            """Resize QEMU resources"""
+            pass
+
+        @resize.command(name="disk")
+        @click.argument("size")
+        def resize_disk(size):
+            """Resize the root disk (e.g., 20G). Run before power on."""
+            self.set_disk_size(size)
+            click.echo(f"Disk will be resized to {size} on next power on")
+
+        @resize.command(name="memory")
+        @click.argument("size")
+        def resize_memory(size):
+            """Set memory size (e.g., 2G, 4G). Takes effect on next boot."""
+            self.set_memory_size(size)
+            click.echo(f"Memory will be set to {size} on next power on")
+
+        return base

packages/jumpstarter-driver-qemu/jumpstarter_driver_qemu/driver.py

@@ -4,6 +4,7 @@
 import logging
 import os
 import platform
+import shutil
 from collections.abc import AsyncGenerator
 from dataclasses import dataclass, field
 from functools import cached_property
@@ -20,7 +21,7 @@
 from jumpstarter_driver_opendal.driver import FlasherInterface
 from jumpstarter_driver_power.driver import PowerInterface, PowerReading
 from jumpstarter_driver_pyserial.driver import PySerial
-from pydantic import BaseModel, Field, validate_call
+from pydantic import BaseModel, ByteSize, Field, TypeAdapter, ValidationError, validate_call
 from qemu.qmp import QMPClient
 from qemu.qmp.protocol import ConnectError, Runstate
 
@@ -169,6 +170,7 @@ async def on(self) -> None:  # noqa: C901
                 proc.check_returncode()
                 info = json.loads(proc.stdout)
                 image_format = info.get("format", "raw")
+                current_virtual_size = info.get("virtual-size") or root.stat().st_size
                 match image_format:
                     case "raw" | "qcow2" | "qcow" | "vmdk":
                         image_driver = image_format
@@ -177,6 +179,34 @@ async def on(self) -> None:  # noqa: C901
             except CalledProcessError:
                 self.logger.warning("unable to detect image format, assuming raw")
                 image_driver = "raw"
+                current_virtual_size = root.stat().st_size
+
+            # Resize disk if configured
+            if self.parent.disk_size:
+                requested = self.parent._parse_size(self.parent.disk_size)
+
+                if requested < current_virtual_size:
+                    raise RuntimeError(
+                        f"Shrinking disk is not supported: current {ByteSize(current_virtual_size).human_readable()}, "
+                        f"requested {self.parent.disk_size}"
+                    )
+
+                available = shutil.disk_usage(root.parent).free
+                if requested > available:
+                    raise RuntimeError(
+                        f"Not enough disk space: need {ByteSize(requested).human_readable()}, "
+                        f"only {ByteSize(available).human_readable()} available"
+                    )
+
+                if requested > current_virtual_size:
+                    self.logger.info(f"Resizing disk to {ByteSize(requested).human_readable()}")
+                    proc = await run_process(
+                        ["qemu-img", "resize", str(root), str(requested)],
+                        stdout=PIPE,
+                        stderr=PIPE,
+                    )
+                    if proc.returncode != 0:
+                        raise RuntimeError(f"Failed to resize disk: {proc.stderr.decode()}")
 
             cmdline += [
                 "-blockdev",
@@ -254,6 +284,7 @@ class Qemu(Driver):
 
     smp: int = 2
     mem: str = "512M"
+    disk_size: str | None = None  # e.g., "20G" (resize disk before boot)
 
     hostname: str = "demo"
     username: str = "jumpstarter"
@@ -372,3 +403,24 @@ def get_username(self) -> str:
     @validate_call(validate_return=True)
     def get_password(self) -> str:
         return self.password
+
+    def _parse_size(self, size: str) -> int:
+        """Parse size string (e.g., '20G') to bytes."""
+        try:
+            return int(TypeAdapter(ByteSize).validate_python(size + "iB" if size[-1] in "kmgtKMGT" else size))
+        except (ValidationError, IndexError):
+            raise ValueError(f"Invalid size: '{size}'. Use e.g. '20G', '512M', '2T'") from None
+
+    @export
+    @validate_call(validate_return=True)
+    def set_disk_size(self, size: str) -> None:
+        """Set the disk size for resizing before boot."""
+        self._parse_size(size)  # Validate
+        self.disk_size = size
+
+    @export
+    @validate_call(validate_return=True)
+    def set_memory_size(self, size: str) -> None:
+        """Set the memory size for next boot."""
+        self._parse_size(size)  # Validate
+        self.mem = size

packages/jumpstarter-driver-qemu/jumpstarter_driver_qemu/driver_test.py

@@ -1,8 +1,11 @@
+import json
 import os
 import platform
 import sys
 import tarfile
 from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, patch
 
 import pytest
 import requests
@@ -13,6 +16,12 @@
 from jumpstarter.common.utils import serve
 
 
+@pytest.fixture
+def anyio_backend():
+    """Use only asyncio backend for anyio tests."""
+    return "asyncio"
+
+
 @pytest.fixture(scope="session")
 def ovmf(tmpdir_factory):
     tmp_path = tmpdir_factory.mktemp("ovmf")
@@ -91,3 +100,104 @@ def test_driver_qemu(tmp_path, ovmf):
             assert s.run("uname -r").stdout.strip() == f"6.11.4-301.fc41.{arch}"
 
         qemu.power.off()
+
+
+@pytest.fixture
+def resize_test():
+    """Create a Qemu driver with a sparse root disk, cleanup after test."""
+    driver = None
+
+    def _create(disk_size, current_size_gb):
+        nonlocal driver
+        driver = Qemu(disk_size=disk_size)
+        root = Path(driver._tmp_dir.name) / "root"
+        root.write_bytes(b"")
+        os.truncate(root, current_size_gb * 1024**3)
+        return driver, current_size_gb * 1024**3
+
+    yield _create
+
+    if driver:
+        driver._tmp_dir.cleanup()
+
+
+def _mock_qemu_img_info(virtual_size):
+    """Return a mock for run_process that simulates qemu-img info."""
+    async def mock(cmd, **kwargs):
+        result = AsyncMock()
+        result.returncode = 0
+        result.stdout = json.dumps({"format": "raw", "virtual-size": virtual_size}).encode()
+        result.check_returncode = lambda: None
+        return result
+    return mock
+
+
+@pytest.mark.anyio
+async def test_resize_shrink_blocked(resize_test):
+    """Shrinking disk should raise RuntimeError."""
+    driver, current = resize_test("10G", 20)  # requested: 10G, current: 20G
+
+    with patch("jumpstarter_driver_qemu.driver.run_process", side_effect=_mock_qemu_img_info(current)):
+        with pytest.raises(RuntimeError, match="Shrinking disk is not supported"):
+            await driver.children["power"].on()
+
+
+@pytest.mark.anyio
+async def test_resize_insufficient_space_blocked(resize_test):
+    """Resize beyond available host space should raise RuntimeError."""
+    driver, current = resize_test("100G", 10)  # requested: 100G, current: 10G
+
+    mock_usage = SimpleNamespace(free=5 * 1024**3)  # only 5G free
+
+    with patch("jumpstarter_driver_qemu.driver.run_process", side_effect=_mock_qemu_img_info(current)):
+        with patch("jumpstarter_driver_qemu.driver.shutil.disk_usage", return_value=mock_usage):
+            with pytest.raises(RuntimeError, match="Not enough disk space"):
+                await driver.children["power"].on()
+
+
+@pytest.mark.anyio
+async def test_resize_succeeds(resize_test):
+    """Resize should call qemu-img resize with correct size."""
+    driver, current = resize_test("20G", 10)  # requested: 20G, current: 10G
+    mock_usage = SimpleNamespace(free=50 * 1024**3)
+
+    with patch("jumpstarter_driver_qemu.driver.run_process", side_effect=_mock_qemu_img_info(current)) as mock_run:
+        with patch("jumpstarter_driver_qemu.driver.shutil.disk_usage", return_value=mock_usage):
+            # Mock Popen to stop before actually starting QEMU VM
+            with patch("jumpstarter_driver_qemu.driver.Popen", side_effect=RuntimeError("mock popen")):
+                with pytest.raises(RuntimeError, match="mock popen"):
+                    await driver.children["power"].on()
+
+    # Find the resize call and verify size argument
+    resize_calls = [c for c in mock_run.call_args_list if "resize" in c.args[0]]
+    assert resize_calls, "qemu-img resize should be called"
+    resize_cmd = resize_calls[0].args[0]  # ['qemu-img', 'resize', path, size]
+    assert resize_cmd[-1] == str(20 * 1024**3)
+
+
+def test_set_disk_size_valid():
+    """Valid size strings should be accepted."""
+    driver = Qemu()
+    driver.set_disk_size("20G")
+    assert driver.disk_size == "20G"
+
+
+def test_set_disk_size_invalid():
+    """Invalid size strings should raise ValueError."""
+    driver = Qemu()
+    with pytest.raises(ValueError, match="Invalid size"):
+        driver.set_disk_size("invalid")
+
+
+def test_set_memory_size_valid():
+    """Valid size strings should be accepted."""
+    driver = Qemu()
+    driver.set_memory_size("2G")
+    assert driver.mem == "2G"
+
+
+def test_set_memory_size_invalid():
+    """Invalid size strings should raise ValueError."""
+    driver = Qemu()
+    with pytest.raises(ValueError, match="Invalid size"):
+        driver.set_memory_size("invalid")