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 ( 

32 patch.object(azure_fileshare, "_share_client") as mock_share_client, 

33 patch.object(mock_share_client, "get_file_client") as mock_get_file_client, 

34 patch.object(mock_share_client, "get_directory_client") as mock_get_directory_client, 

35 ): 

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 ( 

90 patch.object(azure_fileshare, "_share_client") as mock_share_client, 

91 patch.object(mock_share_client, "get_directory_client") as mock_get_directory_client, 

92 patch.object(mock_share_client, "get_file_client") as mock_get_file_client, 

93 ): 

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 ( 

124 patch.object(azure_fileshare, "_share_client") as mock_share_client, 

125 patch.object(mock_share_client, "get_directory_client") as mock_get_directory_client, 

126 patch.object(mock_share_client, "get_file_client") as mock_get_file_client, 

127 ): 

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 ( 

166 patch.object(azure_fileshare, "_share_client") as mock_share_client, 

167 patch.object(mock_share_client, "get_file_client") as mock_get_file_client, 

168 ): 

169 azure_fileshare.upload(config, local_path, remote_path) 

170 

171 mock_get_file_client.assert_called_with(remote_path) 

172 open_path, open_mode = mock_open.call_args.args 

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

174 assert open_mode == "rb" 

175 

176 

177class MyDirEntry: 

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

179 """Dummy class for os.DirEntry.""" 

180 

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

182 self.name = name 

183 self.is_a_dir = is_a_dir 

184 

185 def is_dir(self) -> bool: 

186 return self.is_a_dir 

187 

188 

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

190 return { 

191 local_folder: [ 

192 MyDirEntry("a_folder", True), 

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

194 ], 

195 f"{local_folder}/a_folder": [ 

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

197 ], 

198 } 

199 

200 

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

202 return { 

203 local_folder: True, 

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

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

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

207 } 

208 

209 

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

211 skip_prefix = os.getcwd() 

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

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

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

215 else: 

216 result = input_path 

217 # Change file seps to unix-style 

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

219 

220 

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

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

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

224def test_upload_directory_non_recursive( 

225 mock_scandir: MagicMock, 

226 mock_isdir: MagicMock, 

227 mock_open: MagicMock, 

228 azure_fileshare: AzureFileShareService, 

229) -> None: 

230 remote_folder = "a/remote/folder" 

231 local_folder = "some/local/folder" 

232 scandir_returns = make_scandir_returns(local_folder) 

233 isdir_returns = make_isdir_returns(local_folder) 

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

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

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

237 

238 config: dict = {} 

239 with ( 

240 patch.object(azure_fileshare, "_share_client") as mock_share_client, 

241 patch.object(mock_share_client, "get_file_client") as mock_get_file_client, 

242 ): 

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

244 

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

246 

247 

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

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

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

251def test_upload_directory_recursive( 

252 mock_scandir: MagicMock, 

253 mock_isdir: MagicMock, 

254 mock_open: MagicMock, 

255 azure_fileshare: AzureFileShareService, 

256) -> None: 

257 remote_folder = "a/remote/folder" 

258 local_folder = "some/local/folder" 

259 scandir_returns = make_scandir_returns(local_folder) 

260 isdir_returns = make_isdir_returns(local_folder) 

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

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

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

264 

265 config: dict = {} 

266 with ( 

267 patch.object(azure_fileshare, "_share_client") as mock_share_client, 

268 patch.object(mock_share_client, "get_file_client") as mock_get_file_client, 

269 ): 

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

271 

272 mock_get_file_client.assert_has_calls( 

273 [ 

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

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

276 ], 

277 any_order=True, 

278 )