diff --git a/ext4_symlink.py b/ext4_symlink.py index 86913ab..fd59402 100644 --- a/ext4_symlink.py +++ b/ext4_symlink.py @@ -191,14 +191,57 @@ def _ps_single_quote(s: str) -> str: return s.replace("'", "''") -def _disk_number(vhd_path: str) -> int | None: - """OS disk number of the attached VHD (== PhysicalDriveN index).""" +_ASSOC_DISK_RE = re.compile(r"Associated disk#\s*:\s*(\d+)", re.IGNORECASE) +_ASSOC_DISK_FIELD_RE = re.compile(r"Associated disk#", re.IGNORECASE) + + +def _disk_number_via_diskpart(vhd_path: str) -> tuple[int | None, bool]: + """``(disk_number, parsed)`` from ``diskpart``'s own ``detail vdisk``. + + diskpart already knows which disk it just attached and prints it as + ``Associated disk#: 2``, so asking it costs one diskpart run (~1.6s) where + the ``Get-Disk`` route costs ~17s -- almost all of which is PowerShell cold- + loading the Storage/CIM providers rather than any disk work. That fires + twice per attach/detach cycle and twice per disk per install, so it was the + single largest component of the user's wait. + + ``parsed`` says whether the output was understood at all, which is what lets + the caller fall back safely: the field name is English diskpart text, so a + localised Windows must not silently look like "no disk attached". + """ + r = _diskpart('select vdisk file="%s"\ndetail vdisk\n' % vhd_path) + out = r.stdout or "" + m = _ASSOC_DISK_RE.search(out) + if m: + return int(m.group(1)), True + if _ASSOC_DISK_FIELD_RE.search(out): + return None, True # field present, no number => detached + return None, False # unrecognised output => caller falls back + + +def _disk_number_via_get_disk(vhd_path: str) -> int | None: + """OS disk number via PowerShell ``Get-Disk``. Slow (~17s) but locale-proof.""" r = _run(["powershell", "-NoProfile", "-Command", "(Get-Disk | Where-Object { $_.Location -eq '%s' }).Number" % _ps_single_quote(vhd_path)]) out = (r.stdout or "").strip() return int(out) if out.isdigit() else None +def _disk_number(vhd_path: str) -> int | None: + """OS disk number of the attached VHD (== PhysicalDriveN index), or None. + + Fast path first, ``Get-Disk`` only when diskpart's output isn't recognised. + Both agree in both states -- attached and detached -- which matters because + :func:`_detach` treats ``None`` as proof the disk is gone. + """ + number, parsed = _disk_number_via_diskpart(vhd_path) + if parsed: + return number + logger.debug("diskpart 'detail vdisk' not recognised (localised Windows?); " + "falling back to Get-Disk for %s", vhd_path) + return _disk_number_via_get_disk(vhd_path) + + def _cyg_device(disk_number: int, offset: int) -> str: # Cygwin maps PhysicalDriveN -> /dev/sd; ?offset= selects the ext4 # partition without relying on Cygwin's own partition detection. diff --git a/tests/test_disk_number.py b/tests/test_disk_number.py new file mode 100644 index 0000000..76f1a91 --- /dev/null +++ b/tests/test_disk_number.py @@ -0,0 +1,80 @@ +"""_disk_number: fast diskpart parse, with a Get-Disk fallback for localisation. + +Pins the behaviour that made the fast path safe to adopt -- it must agree with +Get-Disk in BOTH states, because _detach() treats None as proof the disk is gone. +A fast path that returned None for an *attached* disk would make _detach report +success while the image stayed mounted. +""" +from unittest import mock + +import ext4_symlink as es + + +ATTACHED = """ +DiskPart successfully selected the virtual disk file. + +Device type ID: 3 (Unknown) +Virtual size: 128 GB +Physical size: 3530 MB +Filename: C:\\ProgramData\\BlueStacks_nxt\\Engine\\Tiramisu64\\Data.vhdx +Is Child: No +Associated disk#: 2 +""" + +DETACHED = """ +DiskPart successfully selected the virtual disk file. + +Device type ID: 3 (Unknown) +Virtual size: 128 GB +Filename: C:\\ProgramData\\BlueStacks_nxt\\Engine\\Tiramisu64\\Data.vhdx +Is Child: No +Associated disk#: +""" + +# A localised install prints the same data under translated field names. +LOCALISED = """ +DiskPart hat die Datei für den virtuellen Datenträger ausgewählt. + +Datenträger 2 +""" + + +def _completed(stdout): + return mock.Mock(stdout=stdout, stderr="", returncode=0) + + +def test_attached_returns_disk_number_without_calling_get_disk(): + with mock.patch.object(es, "_diskpart", return_value=_completed(ATTACHED)), \ + mock.patch.object(es, "_disk_number_via_get_disk") as slow: + assert es._disk_number("X.vhdx") == 2 + slow.assert_not_called() + + +def test_detached_returns_none_without_calling_get_disk(): + # _detach() relies on this: None must mean "gone", not "couldn't tell". + with mock.patch.object(es, "_diskpart", return_value=_completed(DETACHED)), \ + mock.patch.object(es, "_disk_number_via_get_disk") as slow: + assert es._disk_number("X.vhdx") is None + slow.assert_not_called() + + +def test_unrecognised_output_falls_back_to_get_disk(): + # Must NOT be mistaken for "detached" -- that would let _detach() claim + # success on a still-mounted image on any non-English Windows. + with mock.patch.object(es, "_diskpart", return_value=_completed(LOCALISED)), \ + mock.patch.object(es, "_disk_number_via_get_disk", return_value=7) as slow: + assert es._disk_number("X.vhdx") == 7 + slow.assert_called_once() + + +def test_parse_flag_distinguishes_detached_from_unparseable(): + with mock.patch.object(es, "_diskpart", return_value=_completed(DETACHED)): + assert es._disk_number_via_diskpart("X.vhdx") == (None, True) + with mock.patch.object(es, "_diskpart", return_value=_completed(LOCALISED)): + assert es._disk_number_via_diskpart("X.vhdx") == (None, False) + + +def test_case_insensitive_field_match(): + with mock.patch.object(es, "_diskpart", + return_value=_completed("associated DISK#: 11\n")): + assert es._disk_number_via_diskpart("X.vhdx") == (11, True)