cv4pve-cli-tui/tests/test_pct_shrink.py

"""Tests for scripts/pct-shrink-rootfs; spec: docs/plans/lxc-shrink.md (section B).

The script runs for real, but pct, df, rsync, flock, ha-manager, pvesr and id are
stubs first on PATH. The pct stub edits a temp <vmid>.conf the way PVE does (mp9 set,
mpN -> first free unusedK on delete, lock: mounted on mount); every stub logs its argv.
"""

import os
import subprocess
from pathlib import Path

import pytest

SCRIPT = Path(__file__).resolve().parent.parent / "scripts" / "pct-shrink-rootfs"
GIB = 1024**3

STUBS = {
    "id": 'echo "${STUB_UID:-0}"',
    "flock": 'echo "flock $*" >> "$STUB_LOG"\nshift\nexec "$@"',
    "ha-manager": 'echo "${STUB_HA_OUT:-}"',
    "pvesr": 'echo "${STUB_PVESR_OUT:-}"',
    "df": """echo "df $*" >> "$STUB_LOG"
# like GNU df: -i and --output cannot be combined
if [[ " $* " == *" -i "* && "$*" == *--output* ]]; then echo "df: options -i and --output are mutually exclusive" >&2; exit 1; fi
if [[ "$*" == *--output=iused* ]]; then echo IUsed; echo "${STUB_IUSED:-1000}"
else echo Used; echo "${STUB_USED:-1073741824}"; fi""",
    "rsync": """echo "rsync $*" >> "$STUB_LOG"
[[ -z ${STUB_RSYNC_FAIL:-} ]]""",
    "pct": """echo "pct $*" >> "$STUB_LOG"
conf="$PCT_CONF_DIR/$2.conf"
case "$1" in
  status) echo "status: ${STUB_STATUS:-stopped}" ;;
  config) cat "$conf" ;;
  mount) echo "lock: mounted" >> "$conf" ;;
  unmount)
    n=$(grep -c '^pct unmount' "$STUB_LOG")
    [[ ${STUB_FAIL_UNMOUNT:-0} == "$n" ]] && exit 1
    grep -vx 'lock: mounted' "$conf" > "$conf.new"; mv "$conf.new" "$conf" ;;
  set)
    if [[ $3 == --mp9 ]]; then
      spec=${4%%,*}; storage=${spec%%:*}; gib=${spec#*:}
      vol=${STUB_NEW_VOLID:-$storage:vm-$2-disk-1}
      echo "mp9: $vol,mp=/.shrink-new,size=${gib}G" >> "$conf"
    else
      key=$4
      line=$(grep "^$key: " "$conf")
      grep -v "^$key: " "$conf" > "$conf.new"
      if [[ $key == mp* ]]; then
        vol=${line#*: }; vol=${vol%%,*}; k=0
        while grep -q "^unused$k: " "$conf.new"; do k=$((k+1)); done
        echo "unused$k: $vol" >> "$conf.new"
      fi
      mv "$conf.new" "$conf"
    fi ;;
esac""",
}

BASE = """\
arch: amd64
cores: 2
hostname: web
memory: 512
net0: name=eth0,bridge=vmbr0,ip=dhcp
ostype: debian
rootfs: local-lvm:vm-101-disk-0,size=8G,acl=1
swap: 512
unused0: local-lvm:vm-101-disk-5
unused2: local-lvm:vm-101-disk-7
"""


class Node:
    def __init__(self, tmp: Path):
        self.bin, self.confdir, self.backup = tmp / "bin", tmp / "lxc", tmp / "backup"
        for d in (self.bin, self.confdir, self.backup):
            d.mkdir()
        for name, body in STUBS.items():
            (self.bin / name).write_text(f"#!/usr/bin/env bash\n{body}\n")
            (self.bin / name).chmod(0o755)
        self.conf = self.confdir / "101.conf"
        self.log = tmp / "calls.log"
        self.log.touch()
        self.conf.write_text(BASE)

    def run(self, *args, **env):
        full = {
            **os.environ,
            "PATH": f"{self.bin}:{os.environ['PATH']}",
            "PCT_CONF_DIR": str(self.confdir),
            "PCT_SHRINK_BACKUP_DIR": str(self.backup),
            "STUB_LOG": str(self.log),
            **env,
        }
        return subprocess.run(
            [str(SCRIPT), *args], check=False, env=full, capture_output=True, text=True
        )

    def calls(self, *prefixes):
        return [x for x in self.log.read_text().splitlines() if x.startswith(prefixes)]


@pytest.fixture
def node(tmp_path):
    return Node(tmp_path)


MUTATING = ("pct set", "rsync", "flock")

# name: (extra conf text, replacement for the rootfs line or None, env, stderr keyword)
REFUSALS = {
    "not-root": ("", None, {"STUB_UID": "1000"}, "root"),
    "running": ("", None, {"STUB_STATUS": "running"}, "stopped"),
    "snapshot": ("[snap1]\ncores: 1\n", None, {}, "section"),
    "pending": ("[pve:pending]\ncores: 4\n", None, {}, "section"),
    "lock": ("lock: backup\n", None, {}, "lock"),
    "template": ("template: 1\n", None, {}, "template"),
    "protection": ("protection: 1\n", None, {}, "protection"),
    "ha": ("", None, {"STUB_HA_OUT": "service ct:101 (node1, started)"}, "HA"),
    "replication": (
        "",
        None,
        {"STUB_PVESR_OUT": "101-0 node2 */15 - yes"},
        "replication",
    ),
    "bind-rootfs": ("", "rootfs: /srv/ct101,size=8G", {}, "not a volume"),
    "mp": ("mp0: local-lvm:vm-101-disk-2,mp=/data,size=1G\n", None, {}, "mount point"),
    "mount-entry": ("lxc.mount.entry: /a b none bind 0 0\n", None, {}, "mount point"),
    "too-full": ("", None, {"STUB_USED": str(2 * GIB)}, "too full"),
    "too-many-inodes": ("", None, {"STUB_IUSED": "200000"}, "inodes"),
}


@pytest.mark.parametrize("name", REFUSALS)
def test_refusal_changes_nothing(node, name):
    extra, rootfs, env, keyword = REFUSALS[name]
    text = BASE + extra
    if rootfs:
        text = text.replace("rootfs: local-lvm:vm-101-disk-0,size=8G,acl=1", rootfs)
    node.conf.write_text(text)
    before = node.conf.read_bytes()
    r = node.run("--yes", "101", "2", **env)
    assert r.returncode != 0
    assert keyword in r.stderr
    assert node.conf.read_bytes() == before
    assert node.calls(*MUTATING) == []
    assert list(node.backup.iterdir()) == []


@pytest.mark.parametrize(
    "args",
    [
        ("101", "8"),
        ("101", "9"),
        ("101", "0"),
        ("101", "abc"),
        ("101", "-1"),
        ("101", "1.5"),
        ("101",),
        ("101", "2", "3"),
        ("abc", "2"),
    ],
)
def test_bad_arguments_or_size_not_smaller(node, args):
    before = node.conf.read_bytes()
    r = node.run("--yes", *args)
    assert r.returncode != 0
    assert r.stderr.strip()
    assert node.conf.read_bytes() == before
    assert node.calls(*MUTATING) == []


def test_missing_config_refused(node):
    r = node.run("--yes", "999", "2")
    assert r.returncode != 0
    assert "999.conf" in r.stderr


def test_pct_missing_refused(node):
    (node.bin / "pct").unlink()
    # Minimal PATH, so a real pct on the host cannot be found either.
    r = node.run("--yes", "101", "2", PATH=f"{node.bin}:/usr/bin:/bin")
    assert r.returncode != 0
    assert "pct not found" in r.stderr
    assert node.conf.read_bytes() == BASE.encode()


def test_dry_run_leaves_config_and_prints_plan(node):
    r = node.run("101", "2", STUB_HA_OUT="service ct:1010 (node1, started)")
    assert r.returncode == 0, r.stderr
    assert node.conf.read_bytes() == BASE.encode()
    assert node.calls("pct mount", "pct unmount") == [
        "pct mount 101",
        "pct unmount 101",
    ]
    assert node.calls(*MUTATING) == []
    assert str(node.confdir) in r.stdout
    assert str(node.backup) in r.stdout
    assert "--yes" in r.stdout


def run_happy(node, **env):
    r = node.run("--yes", "101", "2", **env)
    assert r.returncode == 0, r.stderr
    return r


def test_happy_path_exact_config_and_call_sequence(node):
    r = run_happy(node)
    assert (
        node.conf.read_text()
        == BASE.replace(
            "rootfs: local-lvm:vm-101-disk-0,size=8G,acl=1",
            "rootfs: local-lvm:vm-101-disk-1,size=2G,acl=1",
        )
        + "unused1: local-lvm:vm-101-disk-0\n"
    )
    assert node.calls("pct", "rsync") == [
        "pct status 101",
        "pct mount 101",
        "pct unmount 101",
        "pct set 101 --mp9 local-lvm:2,mp=/.shrink-new",
        "pct mount 101",
        (
            "rsync -aHAX --numeric-ids --whole-file --sparse --one-file-system"
            " --exclude=/.shrink-new /var/lib/lxc/101/rootfs/ /var/lib/lxc/101/rootfs/.shrink-new/"
        ),
        "pct unmount 101",
    ]
    assert (
        node.log.read_text()
        .splitlines()[-1]
        .startswith("flock /run/lock/lxc/pve-config-101.lock")
    )
    backups = list(node.backup.glob("pct-shrink-101-*.conf"))
    # Copied under the lock just before the swap, so it already holds mp9.
    mp9 = "mp9: local-lvm:vm-101-disk-1,mp=/.shrink-new,size=2G\n"
    assert [b.read_text() for b in backups] == [BASE + mp9]
    assert "unused1: local-lvm:vm-101-disk-0" in r.stdout
    assert "pct set 101 --delete unused1" in r.stdout
    assert "chattr" in r.stdout


def test_happy_path_dir_style_volid(node):
    node.conf.write_text(
        BASE.replace("local-lvm:vm-101-disk-0", "local:101/vm-101-disk-0.raw")
        .replace("unused0: local-lvm:vm-101-disk-5\n", "")
        .replace("unused2: local-lvm:vm-101-disk-7\n", "")
    )
    run_happy(node, STUB_NEW_VOLID="local:101/vm-101-disk-1.raw")
    assert (
        "rootfs: local:101/vm-101-disk-1.raw,size=2G,acl=1\n" in node.conf.read_text()
    )
    assert node.conf.read_text().endswith("unused0: local:101/vm-101-disk-0.raw\n")
    assert "mp9" not in node.conf.read_text()


def test_rsync_failure_frees_new_volume_by_value(node):
    r = node.run("--yes", "101", "2", STUB_RSYNC_FAIL="1")
    assert r.returncode != 0
    assert node.conf.read_bytes() == BASE.encode()
    # mp9 moved to the first free unusedK (1), which is found by value and deleted
    assert node.calls("pct")[-3:] == [
        "pct unmount 101",
        "pct set 101 --delete mp9",
        "pct set 101 --delete unused1",
    ]
    assert node.calls("flock") == []
    assert list(node.backup.iterdir()) == []


def test_unmount_failure_while_measuring_stops_before_any_set(node):
    r = node.run("--yes", "101", "2", STUB_FAIL_UNMOUNT="1")
    assert r.returncode != 0
    assert "pct unlock 101" in r.stderr
    assert node.calls("pct unmount") == ["pct unmount 101"]
    assert node.calls(*MUTATING) == []


def test_unmount_failure_after_rsync_stops_without_further_set(node):
    r = node.run("--yes", "101", "2", STUB_FAIL_UNMOUNT="2")
    assert r.returncode != 0
    assert "pct unlock 101" in r.stderr
    assert node.calls("pct set") == ["pct set 101 --mp9 local-lvm:2,mp=/.shrink-new"]
    assert node.calls("flock") == []
    text = node.conf.read_text()
    assert "rootfs: local-lvm:vm-101-disk-0,size=8G,acl=1" in text
    assert "mp9: local-lvm:vm-101-disk-1" in text
    assert list(node.backup.iterdir()) == []