furtka/tests/test_drives.py

import json

import drives
import pytest
from drives import (
    MIN_DISK_GB,
    _boot_device_from_cmdline,
    classify_drive,
    format_size,
    get_drive_type_score,
    get_size_score,
    parse_lsblk_json,
)

GB = 1024**3


def _lsblk(*entries):
    """Build lsblk -J output. Each entry: (name, size_gb, type, extra-dict)."""
    devs = []
    for name, size_gb, dev_type, *extra in entries:
        d = {
            "name": name,
            "size": int(size_gb * GB),
            "type": dev_type,
            "ro": False,
            "rm": False,
            "tran": "sata",
            "rota": True,
        }
        if extra:
            d.update(extra[0])
        devs.append(d)
    return json.dumps({"blockdevices": devs})


@pytest.fixture(autouse=True)
def _smart_passed(monkeypatch):
    monkeypatch.setattr(drives, "_smart_status", lambda _: "passed")


def _names(devices):
    return [d["name"] for d in devices]


def test_classify_uses_transport_and_rotation():
    assert classify_drive("nvme0n1", tran="nvme", rota=False) == "NVMe"
    assert classify_drive("sda", tran="sata", rota=False) == "SSD"
    assert classify_drive("sda", tran="sata", rota=True) == "HDD"
    assert classify_drive("mmcblk0", tran="", rota=False) == "eMMC"
    assert classify_drive("sdb", tran="usb", rota=False) == "USB"
    assert classify_drive("sdc", tran="", rota=True, removable=True) == "USB"
    # SD cards in a built-in reader are removable mmc — still eMMC/mmc class.
    assert classify_drive("mmcblk1", tran="", rota=False, removable=True) == "eMMC"


def test_type_scores_rank_internal_over_usb():
    assert get_drive_type_score("NVMe") > get_drive_type_score("SSD")
    assert get_drive_type_score("SSD") > get_drive_type_score("HDD")
    assert get_drive_type_score("HDD") > get_drive_type_score("USB")


def test_size_score_bands():
    assert get_size_score(None) == 5
    assert get_size_score(64) == 5
    assert get_size_score(256) == 7
    assert get_size_score(1024) == 10


def test_format_size():
    assert format_size(500 * GB) == "500 GB"
    assert format_size(1024 * GB) == "1 TB"
    assert format_size(int(1.5 * 1024 * GB)) == "1.5 TB"
    assert format_size(int(7.5 * GB)) == "7.5 GB"


def test_drops_loop_rom_zram_and_emmc_boot_areas():
    out = _lsblk(
        ("loop0", 2.5, "loop"),
        ("sr0", 1, "rom"),
        ("zram0", 4, "disk"),
        ("mmcblk0boot0", 0.004, "disk"),
        ("mmcblk0boot1", 0.004, "disk"),
        ("mmcblk0rpmb", 0.004, "disk"),
        ("mmcblk0", 64, "disk", {"tran": "", "rota": False}),
    )
    assert _names(parse_lsblk_json(out)) == ["/dev/mmcblk0"]


def test_drops_empty_card_reader_slots_and_read_only_disks():
    out = _lsblk(
        ("sdc", 0, "disk", {"rm": True, "tran": "usb"}),
        ("sdd", 500, "disk", {"ro": True}),
        ("sda", 500, "disk"),
    )
    assert _names(parse_lsblk_json(out)) == ["/dev/sda"]


def test_ssd_beats_bigger_hdd():
    # Regression: sdX SSDs used to be labelled HDD, so the 2 TB spinning disk
    # won the recommendation over the 256 GB SSD.
    out = _lsblk(
        ("sda", 2048, "disk", {"rota": True}),
        ("sdb", 256, "disk", {"rota": False}),
    )
    devices = parse_lsblk_json(out)
    assert _names(devices) == ["/dev/sdb", "/dev/sda"]
    assert [d["type_label"] for d in devices] == ["SSD", "HDD"]


def test_usb_disk_ranked_after_internal():
    out = _lsblk(
        ("sda", 32, "disk", {"tran": "usb", "rota": False, "rm": True}),
        ("mmcblk0", 32, "disk", {"tran": "", "rota": False}),
    )
    devices = parse_lsblk_json(out)
    assert _names(devices) == ["/dev/mmcblk0", "/dev/sda"]
    assert devices[1]["warning"]


def test_too_small_disk_listed_but_unusable_and_last():
    out = _lsblk(
        ("sda", MIN_DISK_GB - 1, "disk", {"rota": False}),
        ("sdb", 500, "disk"),
    )
    devices = parse_lsblk_json(out)
    assert _names(devices) == ["/dev/sdb", "/dev/sda"]
    assert devices[0]["usable"] is True
    assert devices[1]["usable"] is False
    assert str(MIN_DISK_GB) in devices[1]["reason"]


def test_bios_warns_about_disks_over_2tb():
    out = _lsblk(("sda", 4096, "disk"))
    [bios] = parse_lsblk_json(out, uefi=False)
    [efi] = parse_lsblk_json(out, uefi=True)
    assert "2 TB" in bios["warning"]
    assert efi["warning"] == ""


def test_surfaces_smart_warning(monkeypatch):
    monkeypatch.setattr(drives, "_smart_status", lambda _: "failed")
    [dev] = parse_lsblk_json(_lsblk(("sda", 500, "disk")))
    assert dev["health_label"] == "SMART warning"


def test_accepts_old_util_linux_string_flags():
    out = json.dumps(
        {
            "blockdevices": [
                {
                    "name": "sda",
                    "size": str(500 * GB),
                    "type": "disk",
                    "ro": "0",
                    "rm": "0",
                    "tran": "sata",
                    "rota": "0",
                },
            ]
        }
    )
    [dev] = parse_lsblk_json(out)
    assert dev["type_label"] == "SSD"


def test_handles_empty_and_garbage_output():
    assert parse_lsblk_json("") == []
    assert parse_lsblk_json("not json") == []


def test_drops_boot_usb():
    out = _lsblk(
        ("sda", 500, "disk"),
        ("sdb", 16, "disk", {"tran": "usb"}),
        ("nvme0n1", 1024, "disk", {"tran": "nvme", "rota": False}),
    )
    assert _names(parse_lsblk_json(out, boot_disk="sdb")) == ["/dev/nvme0n1", "/dev/sda"]


def test_boot_device_from_cmdline(monkeypatch):
    calls = []

    class R:
        returncode = 0
        stdout = "/dev/sdb1\n"

    monkeypatch.setattr(drives, "_run", lambda cmd, **_: calls.append(cmd) or R())
    cmdline = "BOOT_IMAGE=/vmlinuz archisosearchuuid=2026-08-25-10-00-00-00 quiet"
    assert _boot_device_from_cmdline(cmdline) == "/dev/sdb1"
    assert calls[-1] == ["blkid", "-U", "2026-08-25-10-00-00-00"]
    assert _boot_device_from_cmdline("archisolabel=FURTKA_2620") == "/dev/sdb1"
    assert calls[-1] == ["blkid", "-L", "FURTKA_2620"]
    assert _boot_device_from_cmdline("archisodevice=/dev/sr0") == "/dev/sr0"
    assert _boot_device_from_cmdline("quiet splash") is None


def test_boot_disk_whole_disk_iso(monkeypatch):
    # dd-written stick: ISO filesystem on the whole disk, PKNAME empty.
    class R:
        def __init__(self, out):
            self.returncode = 0
            self.stdout = out

    def fake_run(cmd, **_):
        if cmd[0] == "findmnt":
            return R("/dev/sdb\n")
        return R("\n")

    monkeypatch.setattr(drives, "_run", fake_run)
    assert drives._boot_disk_name() == "sdb"