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"