Lufus/tests/test_find_usb.py at f77fd1cd2fb346b8a057bb62841658c185bfa7fa · R-Larocque/Lufus · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
from __future__ import annotations
import sys
import os
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import MagicMock

ROOT = Path(__file__).resolve().parents[1]
SRC = ROOT / "src"
if str(SRC) not in sys.path:
    sys.path.insert(0, str(SRC))

from lufus.drives import find_usb as find_usb_module


def test_find_usb_returns_mount_to_label_mapping(monkeypatch) -> None:
    user = "testuser"
    mount_path = f"/media/{user}/MY_USB"
    device_node = "/dev/sdb1"

    monkeypatch.setattr(find_usb_module.getpass, "getuser", lambda: user)
    monkeypatch.setattr(
        find_usb_module.os.path,
        "exists",
        lambda p: p in {"/media", f"/media/{user}", mount_path},
    )
    monkeypatch.setattr(
        find_usb_module.os.path,
        "isdir",
        lambda p: p in {"/media", f"/media/{user}", mount_path},
    )
    monkeypatch.setattr(
        find_usb_module.os,
        "listdir",
        lambda p: ["MY_USB"] if p == f"/media/{user}" else [],
    )
    monkeypatch.setattr(
        find_usb_module.psutil,
        "disk_partitions",
        lambda *args, **kwargs: [SimpleNamespace(mountpoint=mount_path, device=device_node)],
    )

    # Mock os.stat safely
    os_stat_orig = os.stat
    def mock_os_stat(p):
        if str(p).startswith("/dev/"):
            mock_stat = MagicMock()
            mock_stat.st_rdev = 1234
            return mock_stat
        return os_stat_orig(p)
    monkeypatch.setattr(find_usb_module.os, "stat", mock_os_stat)

    # Mock pyudev
    mock_context = MagicMock()
    mock_device = MagicMock()
    mock_device.get.return_value = "lufus_USB"
    monkeypatch.setattr(find_usb_module.pyudev, "Context", lambda: mock_context)
    monkeypatch.setattr(find_usb_module.pyudev.Devices, "from_device_number", lambda ctx, type, num: mock_device)

    result = find_usb_module.find_usb()
    assert result == {mount_path: "lufus_USB"}


def test_find_usb_falls_back_to_dir_name_when_pyudev_fails(monkeypatch) -> None:
    user = "testuser"
    mount_path = f"/media/{user}/NO_LABEL"
    device_node = "/dev/sdc1"

    monkeypatch.setattr(find_usb_module.getpass, "getuser", lambda: user)
    monkeypatch.setattr(
        find_usb_module.os.path,
        "exists",
        lambda p: p in {"/media", f"/media/{user}", mount_path},
    )
    monkeypatch.setattr(
        find_usb_module.os.path,
        "isdir",
        lambda p: p in {"/media", f"/media/{user}", mount_path},
    )
    monkeypatch.setattr(
        find_usb_module.os,
        "listdir",
        lambda p: ["NO_LABEL"] if p == f"/media/{user}" else [],
    )
    monkeypatch.setattr(
        find_usb_module.psutil,
        "disk_partitions",
        lambda *args, **kwargs: [SimpleNamespace(mountpoint=mount_path, device=device_node)],
    )

    # Mock os.stat safely
    os_stat_orig = os.stat
    def mock_os_stat(p):
        if str(p).startswith("/dev/"):
            mock_stat = MagicMock()
            mock_stat.st_rdev = 5678
            return mock_stat
        return os_stat_orig(p)
    monkeypatch.setattr(find_usb_module.os, "stat", mock_os_stat)

    # Mock pyudev to fail
    mock_context = MagicMock()
    monkeypatch.setattr(find_usb_module.pyudev, "Context", lambda: mock_context)
    monkeypatch.setattr(find_usb_module.pyudev.Devices, "from_device_number", MagicMock(side_effect=Exception("udev fail")))

    result = find_usb_module.find_usb()
    assert result == {mount_path: "NO_LABEL"}


def test_find_dn_returns_matching_device_node(monkeypatch) -> None:
    user = "testuser"
    mount_path = f"/media/{user}/FLASH"

    monkeypatch.setattr(find_usb_module.getpass, "getuser", lambda: user)
    monkeypatch.setattr(
        find_usb_module.os.path,
        "exists",
        lambda p: p in {"/media", f"/media/{user}", mount_path},
    )
    monkeypatch.setattr(
        find_usb_module.os.path,
        "isdir",
        lambda p: p in {"/media", f"/media/{user}", mount_path},
    )
    monkeypatch.setattr(
        find_usb_module.os,
        "listdir",
        lambda p: ["FLASH"] if p == f"/media/{user}" else [],
    )
    monkeypatch.setattr(
        find_usb_module.psutil,
        "disk_partitions",
        lambda *args, **kwargs: [SimpleNamespace(mountpoint=mount_path, device="/dev/sdd1")],
    )

    assert find_usb_module.find_device_node() == "/dev/sdd1"