prole/tests/installer/test_storage_probe.py

from __future__ import annotations

import subprocess

from knoe.core import storage_probe as sp


class _Node:
    def __init__(self, name: str, filesystems: list[dict]):
        self.name = name
        self.filesystems = filesystems


def _mk_result(rr: float | None, rw: float | None, lat: float | None) -> sp.StorageProbeResult:
    return sp.StorageProbeResult(
        target_id="t1",
        node_name="node-a",
        label="Disk",
        mountpoint="/data",
        source_type="filesystem_mount",
        storage_class="ssd",
        succeeded=True,
        skipped=False,
        error=None,
        metrics=sp.StorageProbeMetrics(
            randread_iops=rr,
            randwrite_iops=rw,
            read_lat_ms_p95=lat,
            write_lat_ms_p95=lat,
        ),
        io_class=None,
        timing_multiplier=None,
        observed_at="2026-01-01T00:00:00+00:00",
    )


def test_discover_storage_probe_targets_dedupes_and_labels() -> None:
    node = _Node(
        "merlin",
        [
            {
                "mountpoint": "/synology/d004",
                "device": "/dev/sdb1",
                "fstype": "ext4",
                "storageClass": "synology",
                "transport": "sata",
                "rotational": True,
            },
            {
                "mountpoint": "/synology/d004",
                "device": "/dev/sdb1",
                "fstype": "ext4",
                "storageClass": "synology",
            },
            {
                "mountpoint": "/var/lib",
                "device": "/dev/nvme0n1p1",
                "fstype": "xfs",
                "storageClass": "nvme",
                "transport": "nvme",
                "rotational": False,
            },
        ],
    )

    targets = sp.discover_storage_probe_targets(node)

    assert len([t for t in targets if t.mountpoint == "/synology/d004"]) == 1
    synology = next(t for t in targets if t.mountpoint == "/synology/d004")
    assert synology.label == "Synology d004"
    varlib = next(t for t in targets if t.mountpoint == "/var/lib")
    assert "filesystem" in varlib.label.lower()


def test_parse_fio_json_metrics() -> None:
    payload = {
        "jobs": [
            {
                "jobname": "randread",
                "read": {
                    "iops": 620.5,
                    "bw_bytes": 2541568,
                    "clat_ns": {"percentile": {"95.000000": 24700000}},
                },
                "write": {},
            },
            {
                "jobname": "randwrite",
                "read": {},
                "write": {
                    "iops": 240.2,
                    "bw_bytes": 983040,
                    "clat_ns": {"percentile": {"95.000000": 19500000}},
                },
            },
        ]
    }

    metrics = sp.parse_fio_json_metrics(payload)

    assert metrics.randread_iops == 620.5
    assert metrics.randwrite_iops == 240.2
    assert metrics.read_bw_bytes == 2541568
    assert metrics.write_bw_bytes == 983040
    assert metrics.read_lat_ms_p95 == 24.7
    assert metrics.write_lat_ms_p95 == 19.5


def test_run_storage_probe_falls_back_when_fio_missing(tmp_path, monkeypatch) -> None:
    target_dir = tmp_path / "probe"
    target_dir.mkdir(parents=True)
    target = sp.StorageProbeTarget(
        node_name="node-a",
        source_type="filesystem_mount",
        mountpoint=str(target_dir),
        device=None,
        filesystem="ext4",
        storage_class="ssd",
        label="Local SSD mount",
        path_exists=True,
        writable=True,
        notes=[],
    )

    monkeypatch.setattr(sp.shutil, "which", lambda _tool: None)

    def fake_run(cmd, capture_output, text, timeout):
        if cmd[0] != "dd":
            raise AssertionError(f"unexpected command: {cmd}")
        return subprocess.CompletedProcess(
            args=cmd,
            returncode=0,
            stdout="",
            stderr="268435456 bytes transferred in 0.50 s",
        )

    monkeypatch.setattr(sp.subprocess, "run", fake_run)

    result = sp.run_storage_probe(target, timeout_sec=5, quick=True)

    assert result.succeeded is True
    assert result.approximate is True
    assert result.metrics.randread_iops is not None
    assert result.metrics.randwrite_iops is not None
    assert result.io_class in {"io_cold", "io_slow_durable", "io_general", "io_fast_local"}


def test_classify_storage_io_thresholds() -> None:
    fast = _mk_result(9000, 4500, 2.0)
    assert sp.classify_storage_io(fast)[0] == "io_fast_local"

    general = _mk_result(2500, 1200, 8.0)
    assert sp.classify_storage_io(general)[0] == "io_general"

    slow = _mk_result(600, 220, 20.0)
    assert sp.classify_storage_io(slow)[0] == "io_slow_durable"

    cold = _mk_result(100, 80, 80.0)
    assert sp.classify_storage_io(cold)[0] == "io_cold"


def test_probe_node_storage_cache_hit_and_miss(tmp_path, monkeypatch) -> None:
    mountpoint = tmp_path / "data"
    mountpoint.mkdir(parents=True)
    node = _Node(
        "cache-node",
        [
            {
                "mountpoint": str(mountpoint),
                "device": "/dev/sda1",
                "fstype": "ext4",
                "storageClass": "ssd",
                "transport": "sata",
                "rotational": False,
            }
        ],
    )
    calls = {"n": 0}

    fixed_target = sp.StorageProbeTarget(
        node_name="cache-node",
        source_type="filesystem_mount",
        mountpoint=str(mountpoint),
        device="/dev/sda1",
        filesystem="ext4",
        storage_class="ssd",
        label="Local SSD mount",
        path_exists=True,
        writable=True,
        notes=[],
    )

    monkeypatch.setattr(sp, "discover_storage_probe_targets", lambda _node, _ctx=None: [fixed_target])

    def fake_run(target, timeout_sec, quick):
        calls["n"] += 1
        return sp.StorageProbeResult(
            target_id="abc",
            node_name=target.node_name,
            label=target.label,
            mountpoint=target.mountpoint,
            source_type=target.source_type,
            storage_class=target.storage_class,
            succeeded=True,
            skipped=False,
            error=None,
            metrics=sp.StorageProbeMetrics(randread_iops=1000.0, randwrite_iops=500.0),
            io_class="io_slow_durable",
            timing_multiplier=3.0,
            observed_at="2026-03-26T00:00:00+00:00",
        )

    monkeypatch.setattr(sp, "run_storage_probe", fake_run)

    cache_file = tmp_path / "cache" / "storage_probe_cache.json"
    inv1 = sp.probe_node_storage(node, cache_path=cache_file, quick=True)
    inv2 = sp.probe_node_storage(node, cache_path=cache_file, quick=True)

    assert calls["n"] == 1
    assert len(inv1.items) == 1
    assert len(inv2.items) == 1
    assert inv2.items[0].cached is True