Coverage for mlos_bench/mlos_bench/tests/services/remote/azure/azure_fileshare_test.py: 99%

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5"""Tests for mlos_bench.services.remote.azure.azure_fileshare.""" 

6 

7import os 

8from unittest.mock import MagicMock, Mock, call, patch 

9 

10from mlos_bench.services.remote.azure.azure_fileshare import AzureFileShareService 

11 

12# pylint: disable=missing-function-docstring 

13# pylint: disable=too-many-arguments 

14# pylint: disable=unused-argument 

15 

16 

17@patch("mlos_bench.services.remote.azure.azure_fileshare.open") 

18@patch("mlos_bench.services.remote.azure.azure_fileshare.os.makedirs") 

19def test_download_file( 

20 mock_makedirs: MagicMock, 

21 mock_open: MagicMock, 

22 azure_fileshare: AzureFileShareService, 

23) -> None: 

24 filename = "test.csv" 

25 remote_folder = "a/remote/folder" 

26 local_folder = "some/local/folder" 

27 remote_path = f"{remote_folder}/{filename}" 

28 local_path = f"{local_folder}/{filename}" 

29 

30 config: dict = {} 

31 with patch.object(azure_fileshare, "_share_client") as mock_share_client, patch.object( 

32 mock_share_client, "get_file_client" 

33 ) as mock_get_file_client, patch.object( 

34 mock_share_client, "get_directory_client" 

35 ) as mock_get_directory_client: 

36 

37 mock_get_directory_client.return_value = Mock(exists=Mock(return_value=False)) 

38 

39 azure_fileshare.download(config, remote_path, local_path) 

40 

41 mock_get_file_client.assert_called_with(remote_path) 

42 

43 mock_makedirs.assert_called_with( 

44 local_folder, 

45 exist_ok=True, 

46 ) 

47 open_path, open_mode = mock_open.call_args.args 

48 assert os.path.abspath(local_path) == os.path.abspath(open_path) 

49 assert open_mode == "wb" 

50 

51 

52def make_dir_client_returns(remote_folder: str) -> dict: 

53 return { 

54 remote_folder: Mock( 

55 exists=Mock(return_value=True), 

56 list_directories_and_files=Mock( 

57 return_value=[ 

58 {"name": "a_folder", "is_directory": True}, 

59 {"name": "a_file_1.csv", "is_directory": False}, 

60 ] 

61 ), 

62 ), 

63 f"{remote_folder}/a_folder": Mock( 

64 exists=Mock(return_value=True), 

65 list_directories_and_files=Mock( 

66 return_value=[ 

67 {"name": "a_file_2.csv", "is_directory": False}, 

68 ] 

69 ), 

70 ), 

71 f"{remote_folder}/a_file_1.csv": Mock(exists=Mock(return_value=False)), 

72 f"{remote_folder}/a_folder/a_file_2.csv": Mock(exists=Mock(return_value=False)), 

73 } 

74 

75 

76@patch("mlos_bench.services.remote.azure.azure_fileshare.open") 

77@patch("mlos_bench.services.remote.azure.azure_fileshare.os.makedirs") 

78def test_download_folder_non_recursive( 

79 mock_makedirs: MagicMock, 

80 mock_open: MagicMock, 

81 azure_fileshare: AzureFileShareService, 

82) -> None: 

83 remote_folder = "a/remote/folder" 

84 local_folder = "some/local/folder" 

85 dir_client_returns = make_dir_client_returns(remote_folder) 

86 mock_share_client = azure_fileshare._share_client # pylint: disable=protected-access 

87 

88 config: dict = {} 

89 with patch.object(azure_fileshare, "_share_client") as mock_share_client, patch.object( 

90 mock_share_client, "get_directory_client" 

91 ) as mock_get_directory_client, patch.object( 

92 mock_share_client, "get_file_client" 

93 ) as mock_get_file_client: 

94 

95 mock_get_directory_client.side_effect = lambda x: dir_client_returns[x] 

96 

97 azure_fileshare.download(config, remote_folder, local_folder, recursive=False) 

98 

99 mock_get_file_client.assert_called_with( 

100 f"{remote_folder}/a_file_1.csv", 

101 ) 

102 mock_get_directory_client.assert_has_calls( 

103 [ 

104 call(remote_folder), 

105 call(f"{remote_folder}/a_file_1.csv"), 

106 ], 

107 any_order=True, 

108 ) 

109 

110 

111@patch("mlos_bench.services.remote.azure.azure_fileshare.open") 

112@patch("mlos_bench.services.remote.azure.azure_fileshare.os.makedirs") 

113def test_download_folder_recursive( 

114 mock_makedirs: MagicMock, 

115 mock_open: MagicMock, 

116 azure_fileshare: AzureFileShareService, 

117) -> None: 

118 remote_folder = "a/remote/folder" 

119 local_folder = "some/local/folder" 

120 dir_client_returns = make_dir_client_returns(remote_folder) 

121 

122 config: dict = {} 

123 with patch.object(azure_fileshare, "_share_client") as mock_share_client, patch.object( 

124 mock_share_client, "get_directory_client" 

125 ) as mock_get_directory_client, patch.object( 

126 mock_share_client, "get_file_client" 

127 ) as mock_get_file_client: 

128 mock_get_directory_client.side_effect = lambda x: dir_client_returns[x] 

129 azure_fileshare.download(config, remote_folder, local_folder, recursive=True) 

130 

131 mock_get_file_client.assert_has_calls( 

132 [ 

133 call(f"{remote_folder}/a_file_1.csv"), 

134 call(f"{remote_folder}/a_folder/a_file_2.csv"), 

135 ], 

136 any_order=True, 

137 ) 

138 mock_get_directory_client.assert_has_calls( 

139 [ 

140 call(remote_folder), 

141 call(f"{remote_folder}/a_file_1.csv"), 

142 call(f"{remote_folder}/a_folder"), 

143 call(f"{remote_folder}/a_folder/a_file_2.csv"), 

144 ], 

145 any_order=True, 

146 ) 

147 

148 

149@patch("mlos_bench.services.remote.azure.azure_fileshare.open") 

150@patch("mlos_bench.services.remote.azure.azure_fileshare.os.path.isdir") 

151def test_upload_file( 

152 mock_isdir: MagicMock, 

153 mock_open: MagicMock, 

154 azure_fileshare: AzureFileShareService, 

155) -> None: 

156 filename = "test.csv" 

157 remote_folder = "a/remote/folder" 

158 local_folder = "some/local/folder" 

159 remote_path = f"{remote_folder}/{filename}" 

160 local_path = f"{local_folder}/{filename}" 

161 mock_share_client = azure_fileshare._share_client # pylint: disable=protected-access 

162 mock_isdir.return_value = False 

163 

164 config: dict = {} 

165 with patch.object(azure_fileshare, "_share_client") as mock_share_client, patch.object( 

166 mock_share_client, "get_file_client" 

167 ) as mock_get_file_client: 

168 azure_fileshare.upload(config, local_path, remote_path) 

169 

170 mock_get_file_client.assert_called_with(remote_path) 

171 open_path, open_mode = mock_open.call_args.args 

172 assert os.path.abspath(local_path) == os.path.abspath(open_path) 

173 assert open_mode == "rb" 

174 

175 

176class MyDirEntry: 

177 # pylint: disable=too-few-public-methods 

178 """Dummy class for os.DirEntry.""" 

179 

180 def __init__(self, name: str, is_a_dir: bool): 

181 self.name = name 

182 self.is_a_dir = is_a_dir 

183 

184 def is_dir(self) -> bool: 

185 return self.is_a_dir 

186 

187 

188def make_scandir_returns(local_folder: str) -> dict: 

189 return { 

190 local_folder: [ 

191 MyDirEntry("a_folder", True), 

192 MyDirEntry("a_file_1.csv", False), 

193 ], 

194 f"{local_folder}/a_folder": [ 

195 MyDirEntry("a_file_2.csv", False), 

196 ], 

197 } 

198 

199 

200def make_isdir_returns(local_folder: str) -> dict: 

201 return { 

202 local_folder: True, 

203 f"{local_folder}/a_file_1.csv": False, 

204 f"{local_folder}/a_folder": True, 

205 f"{local_folder}/a_folder/a_file_2.csv": False, 

206 } 

207 

208 

209def process_paths(input_path: str) -> str: 

210 skip_prefix = os.getcwd() 

211 # Remove prefix from os.path.abspath if there 

212 if input_path == os.path.abspath(input_path): 

213 result = input_path[(len(skip_prefix) + 1) :] 

214 else: 

215 result = input_path 

216 # Change file seps to unix-style 

217 return result.replace("\\", "/") 

218 

219 

220@patch("mlos_bench.services.remote.azure.azure_fileshare.open") 

221@patch("mlos_bench.services.remote.azure.azure_fileshare.os.path.isdir") 

222@patch("mlos_bench.services.remote.azure.azure_fileshare.os.scandir") 

223def test_upload_directory_non_recursive( 

224 mock_scandir: MagicMock, 

225 mock_isdir: MagicMock, 

226 mock_open: MagicMock, 

227 azure_fileshare: AzureFileShareService, 

228) -> None: 

229 remote_folder = "a/remote/folder" 

230 local_folder = "some/local/folder" 

231 scandir_returns = make_scandir_returns(local_folder) 

232 isdir_returns = make_isdir_returns(local_folder) 

233 mock_scandir.side_effect = lambda x: scandir_returns[process_paths(x)] 

234 mock_isdir.side_effect = lambda x: isdir_returns[process_paths(x)] 

235 mock_share_client = azure_fileshare._share_client # pylint: disable=protected-access 

236 

237 config: dict = {} 

238 with patch.object(azure_fileshare, "_share_client") as mock_share_client, patch.object( 

239 mock_share_client, "get_file_client" 

240 ) as mock_get_file_client: 

241 azure_fileshare.upload(config, local_folder, remote_folder, recursive=False) 

242 

243 mock_get_file_client.assert_called_with(f"{remote_folder}/a_file_1.csv") 

244 

245 

246@patch("mlos_bench.services.remote.azure.azure_fileshare.open") 

247@patch("mlos_bench.services.remote.azure.azure_fileshare.os.path.isdir") 

248@patch("mlos_bench.services.remote.azure.azure_fileshare.os.scandir") 

249def test_upload_directory_recursive( 

250 mock_scandir: MagicMock, 

251 mock_isdir: MagicMock, 

252 mock_open: MagicMock, 

253 azure_fileshare: AzureFileShareService, 

254) -> None: 

255 remote_folder = "a/remote/folder" 

256 local_folder = "some/local/folder" 

257 scandir_returns = make_scandir_returns(local_folder) 

258 isdir_returns = make_isdir_returns(local_folder) 

259 mock_scandir.side_effect = lambda x: scandir_returns[process_paths(x)] 

260 mock_isdir.side_effect = lambda x: isdir_returns[process_paths(x)] 

261 mock_share_client = azure_fileshare._share_client # pylint: disable=protected-access 

262 

263 config: dict = {} 

264 with patch.object(azure_fileshare, "_share_client") as mock_share_client, patch.object( 

265 mock_share_client, "get_file_client" 

266 ) as mock_get_file_client: 

267 azure_fileshare.upload(config, local_folder, remote_folder, recursive=True) 

268 

269 mock_get_file_client.assert_has_calls( 

270 [ 

271 call(f"{remote_folder}/a_file_1.csv"), 

272 call(f"{remote_folder}/a_folder/a_file_2.csv"), 

273 ], 

274 any_order=True, 

275 )