|
| 1 | +"""Unit tests for sky/utils/cluster_utils.py.""" |
| 2 | +from unittest import mock |
| 3 | + |
| 4 | +import pytest |
| 5 | + |
| 6 | +from sky.utils import cluster_utils |
| 7 | + |
| 8 | + |
| 9 | +class TestConvertWindowsPathToWsl: |
| 10 | + """Tests for _convert_windows_path_to_wsl().""" |
| 11 | + |
| 12 | + def test_standard_windows_path(self): |
| 13 | + result = cluster_utils._convert_windows_path_to_wsl('C:\\Users\\test') |
| 14 | + assert result == '/mnt/c/Users/test' |
| 15 | + |
| 16 | + def test_windows_path_with_forward_slashes(self): |
| 17 | + result = cluster_utils._convert_windows_path_to_wsl('C:/Users/test') |
| 18 | + assert result == '/mnt/c/Users/test' |
| 19 | + |
| 20 | + def test_drive_letter_case_insensitive(self): |
| 21 | + result = cluster_utils._convert_windows_path_to_wsl('D:\\Data') |
| 22 | + assert result == '/mnt/d/Data' |
| 23 | + |
| 24 | + def test_root_drive(self): |
| 25 | + result = cluster_utils._convert_windows_path_to_wsl('C:\\') |
| 26 | + assert result == '/mnt/c/' |
| 27 | + |
| 28 | + def test_already_wsl_path(self): |
| 29 | + result = cluster_utils._convert_windows_path_to_wsl('/already/wsl/path') |
| 30 | + assert result == '/already/wsl/path' |
| 31 | + |
| 32 | + def test_relative_path(self): |
| 33 | + result = cluster_utils._convert_windows_path_to_wsl('relative/path') |
| 34 | + assert result == 'relative/path' |
| 35 | + |
| 36 | + def test_empty_string(self): |
| 37 | + result = cluster_utils._convert_windows_path_to_wsl('') |
| 38 | + assert result == '' |
| 39 | + |
| 40 | + |
| 41 | +class TestConvertWslPathToWindows: |
| 42 | + """Tests for _convert_wsl_path_to_windows().""" |
| 43 | + |
| 44 | + def test_standard_wsl_path(self): |
| 45 | + result = cluster_utils._convert_wsl_path_to_windows('/mnt/c/Users/test') |
| 46 | + assert result == 'C:/Users/test' |
| 47 | + |
| 48 | + def test_different_drive(self): |
| 49 | + result = cluster_utils._convert_wsl_path_to_windows('/mnt/d/Data') |
| 50 | + assert result == 'D:/Data' |
| 51 | + |
| 52 | + def test_root_drive(self): |
| 53 | + result = cluster_utils._convert_wsl_path_to_windows('/mnt/c/') |
| 54 | + assert result == 'C:/' |
| 55 | + |
| 56 | + def test_non_mnt_path(self): |
| 57 | + result = cluster_utils._convert_wsl_path_to_windows('/home/user') |
| 58 | + assert result == '/home/user' |
| 59 | + |
| 60 | + def test_short_mnt_path(self): |
| 61 | + result = cluster_utils._convert_wsl_path_to_windows('/mnt/') |
| 62 | + assert result == '/mnt/' |
| 63 | + |
| 64 | + def test_empty_string(self): |
| 65 | + result = cluster_utils._convert_wsl_path_to_windows('') |
| 66 | + assert result == '' |
| 67 | + |
| 68 | + |
| 69 | +class TestConvertProxyCommandForWindows: |
| 70 | + """Tests for SSHConfigHelper._convert_proxy_command_for_windows().""" |
| 71 | + |
| 72 | + def test_simple_proxy_command(self): |
| 73 | + cmd = 'ssh -W %h:%p user@host' |
| 74 | + result = cluster_utils.SSHConfigHelper._convert_proxy_command_for_windows( |
| 75 | + cmd) |
| 76 | + assert result == "wsl.exe bash -c 'ssh -W %h:%p user@host'" |
| 77 | + |
| 78 | + def test_proxy_command_with_double_quotes(self): |
| 79 | + cmd = 'ssh -o "StrictHostKeyChecking=no" -W %h:%p user@host' |
| 80 | + result = cluster_utils.SSHConfigHelper._convert_proxy_command_for_windows( |
| 81 | + cmd) |
| 82 | + assert result == ( |
| 83 | + "wsl.exe bash -c 'ssh -o \"StrictHostKeyChecking=no\" -W %h:%p user@host'" |
| 84 | + ) |
| 85 | + |
| 86 | + def test_proxy_command_with_single_quotes(self): |
| 87 | + cmd = "ssh -o 'StrictHostKeyChecking=no' -W %h:%p user@host" |
| 88 | + result = cluster_utils.SSHConfigHelper._convert_proxy_command_for_windows( |
| 89 | + cmd) |
| 90 | + # Single quotes should be escaped with '"'"' |
| 91 | + assert result == ( |
| 92 | + "wsl.exe bash -c 'ssh -o '\"'\"'StrictHostKeyChecking=no'\"'\"' -W %h:%p user@host'" |
| 93 | + ) |
| 94 | + |
| 95 | + def test_proxy_command_with_dollar_sign(self): |
| 96 | + # Dollar signs should not be expanded in single-quoted strings |
| 97 | + cmd = 'echo $HOME' |
| 98 | + result = cluster_utils.SSHConfigHelper._convert_proxy_command_for_windows( |
| 99 | + cmd) |
| 100 | + assert result == "wsl.exe bash -c 'echo $HOME'" |
| 101 | + |
| 102 | + def test_proxy_command_with_backticks(self): |
| 103 | + # Backticks should not be expanded in single-quoted strings |
| 104 | + cmd = 'echo `hostname`' |
| 105 | + result = cluster_utils.SSHConfigHelper._convert_proxy_command_for_windows( |
| 106 | + cmd) |
| 107 | + assert result == "wsl.exe bash -c 'echo `hostname`'" |
| 108 | + |
| 109 | + |
| 110 | +class TestGetWslWindowsHome: |
| 111 | + """Tests for get_wsl_windows_home().""" |
| 112 | + |
| 113 | + def test_not_wsl_returns_none(self): |
| 114 | + # Clear the lru_cache to ensure fresh state |
| 115 | + cluster_utils.get_wsl_windows_home.cache_clear() |
| 116 | + with mock.patch('sky.utils.common_utils.is_wsl', return_value=False): |
| 117 | + result = cluster_utils.get_wsl_windows_home() |
| 118 | + assert result is None |
| 119 | + |
| 120 | + def test_wsl_with_userprofile_env(self): |
| 121 | + cluster_utils.get_wsl_windows_home.cache_clear() |
| 122 | + with mock.patch('sky.utils.common_utils.is_wsl', return_value=True): |
| 123 | + with mock.patch.dict('os.environ', |
| 124 | + {'USERPROFILE': 'C:\\Users\\testuser'}): |
| 125 | + with mock.patch('os.path.isdir', return_value=True): |
| 126 | + result = cluster_utils.get_wsl_windows_home() |
| 127 | + assert result == '/mnt/c/Users/testuser' |
| 128 | + |
| 129 | + def test_wsl_without_userprofile_uses_cmd(self): |
| 130 | + cluster_utils.get_wsl_windows_home.cache_clear() |
| 131 | + with mock.patch('sky.utils.common_utils.is_wsl', return_value=True): |
| 132 | + with mock.patch.dict('os.environ', {}, clear=True): |
| 133 | + with mock.patch( |
| 134 | + 'sky.utils.cluster_utils._get_windows_userprofile_via_cmd', |
| 135 | + return_value='D:\\Users\\cmduser'): |
| 136 | + with mock.patch('os.path.isdir', return_value=True): |
| 137 | + result = cluster_utils.get_wsl_windows_home() |
| 138 | + assert result == '/mnt/d/Users/cmduser' |
| 139 | + |
| 140 | + def test_wsl_with_invalid_home_returns_none(self): |
| 141 | + cluster_utils.get_wsl_windows_home.cache_clear() |
| 142 | + with mock.patch('sky.utils.common_utils.is_wsl', return_value=True): |
| 143 | + with mock.patch.dict('os.environ', |
| 144 | + {'USERPROFILE': 'C:\\Users\\testuser'}): |
| 145 | + with mock.patch('os.path.isdir', return_value=False): |
| 146 | + result = cluster_utils.get_wsl_windows_home() |
| 147 | + assert result is None |
| 148 | + |
| 149 | + def test_caching_behavior(self): |
| 150 | + """Test that the function result is cached.""" |
| 151 | + cluster_utils.get_wsl_windows_home.cache_clear() |
| 152 | + call_count = 0 |
| 153 | + |
| 154 | + def mock_is_wsl(): |
| 155 | + nonlocal call_count |
| 156 | + call_count += 1 |
| 157 | + return True |
| 158 | + |
| 159 | + with mock.patch('sky.utils.common_utils.is_wsl', |
| 160 | + side_effect=mock_is_wsl): |
| 161 | + with mock.patch.dict('os.environ', |
| 162 | + {'USERPROFILE': 'C:\\Users\\cached'}): |
| 163 | + with mock.patch('os.path.isdir', return_value=True): |
| 164 | + # Call multiple times |
| 165 | + result1 = cluster_utils.get_wsl_windows_home() |
| 166 | + result2 = cluster_utils.get_wsl_windows_home() |
| 167 | + result3 = cluster_utils.get_wsl_windows_home() |
| 168 | + |
| 169 | + # Should only call is_wsl once due to caching |
| 170 | + assert call_count == 1 |
| 171 | + assert result1 == result2 == result3 == '/mnt/c/Users/cached' |
| 172 | + |
| 173 | + |
| 174 | +class TestGetWindowsUserprofileViaCmd: |
| 175 | + """Tests for _get_windows_userprofile_via_cmd().""" |
| 176 | + |
| 177 | + def test_successful_cmd_call(self): |
| 178 | + mock_result = mock.Mock() |
| 179 | + mock_result.returncode = 0 |
| 180 | + mock_result.stdout = 'C:\\Users\\testuser\n' |
| 181 | + |
| 182 | + with mock.patch('subprocess.run', return_value=mock_result): |
| 183 | + result = cluster_utils._get_windows_userprofile_via_cmd() |
| 184 | + assert result == 'C:\\Users\\testuser' |
| 185 | + |
| 186 | + def test_failed_cmd_call(self): |
| 187 | + mock_result = mock.Mock() |
| 188 | + mock_result.returncode = 1 |
| 189 | + mock_result.stdout = '' |
| 190 | + |
| 191 | + with mock.patch('subprocess.run', return_value=mock_result): |
| 192 | + result = cluster_utils._get_windows_userprofile_via_cmd() |
| 193 | + assert result is None |
| 194 | + |
| 195 | + def test_unexpanded_variable(self): |
| 196 | + mock_result = mock.Mock() |
| 197 | + mock_result.returncode = 0 |
| 198 | + mock_result.stdout = '%USERPROFILE%\n' |
| 199 | + |
| 200 | + with mock.patch('subprocess.run', return_value=mock_result): |
| 201 | + result = cluster_utils._get_windows_userprofile_via_cmd() |
| 202 | + assert result is None |
| 203 | + |
| 204 | + def test_timeout_exception(self): |
| 205 | + import subprocess |
| 206 | + with mock.patch('subprocess.run', |
| 207 | + side_effect=subprocess.TimeoutExpired(cmd='cmd.exe', |
| 208 | + timeout=5)): |
| 209 | + result = cluster_utils._get_windows_userprofile_via_cmd() |
| 210 | + assert result is None |
| 211 | + |
| 212 | + def test_file_not_found(self): |
| 213 | + with mock.patch('subprocess.run', side_effect=FileNotFoundError()): |
| 214 | + result = cluster_utils._get_windows_userprofile_via_cmd() |
| 215 | + assert result is None |
0 commit comments