def test_download_artifacts(ftp_mock): repo = FTPArtifactRepository("ftp://test_ftp/some/path") repo.get_ftp_client = MagicMock() call_mock = MagicMock(return_value=ftp_mock) repo.get_ftp_client.return_value = MagicMock(__enter__=call_mock) # mocked file structure # |- model # |- model.pb # |- variables # |- test.txt dir_name = 'model' # list artifacts at sub directory level file_path = 'model.pb' subdir_name = 'variables' subfile_name = 'test.txt' ftp_dir_name = '/some/path/' + dir_name ftp_mock.cwd = MagicMock(side_effect=[None, None, ftplib.error_perm, None, None, ftplib.error_perm]) ftp_mock.nlst = MagicMock(side_effect=[[file_path, subdir_name], [subfile_name]]) repo.download_artifacts("model") cwd_call_args = sorted([ftp_mock.cwd.call_args_list[i][0][0] for i in range(6)]) expect_args = [ftp_dir_name, ftp_dir_name, os.path.join(ftp_dir_name, file_path), os.path.join(ftp_dir_name, subdir_name), os.path.join(ftp_dir_name, subdir_name), os.path.join(ftp_dir_name, subdir_name, subfile_name)] assert cwd_call_args == expect_args assert ftp_mock.nlst.call_count == 2 assert ftp_mock.retrbinary.call_args_list[0][0][0] == 'RETR ' + file_path assert ftp_mock.retrbinary.call_args_list[1][0][0] == 'RETR ' + subfile_name
def test_download_artifacts(ftp_mock): artifact_root_path = "/some/path" repo = FTPArtifactRepository("ftp://test_ftp" + artifact_root_path) repo.get_ftp_client = MagicMock() call_mock = MagicMock(return_value=ftp_mock) repo.get_ftp_client.return_value = MagicMock(__enter__=call_mock) # mocked file structure # |- model # |- model.pb # |- variables # |- test.txt dir_path = os.path.join(artifact_root_path, 'model') # list artifacts at sub directory level model_file_path_sub = 'model.pb' model_file_path_full = os.path.join(dir_path, model_file_path_sub) subdir_name = 'variables' subdir_path_full = os.path.join(dir_path, subdir_name) subfile_name = 'test.txt' subfile_path_full = os.path.join(artifact_root_path, subdir_path_full, subfile_name) is_dir_mapping = { dir_path: True, model_file_path_full: False, subdir_path_full: True, subfile_path_full: False, } is_dir_call_args = [ dir_path, model_file_path_full, subdir_path_full, model_file_path_full, subdir_path_full, subfile_path_full, subfile_path_full ] cwd_side_effect = [ None if is_dir_mapping[call_arg] else ftplib.error_perm for call_arg in is_dir_call_args ] ftp_mock.cwd = MagicMock(side_effect=cwd_side_effect) ftp_mock.nlst = MagicMock( side_effect=[[model_file_path_sub, subdir_name], [subfile_name]]) repo.download_artifacts("model") cwd_call_args = [ arg_entry[0][0] for arg_entry in ftp_mock.cwd.call_args_list ] assert cwd_call_args == is_dir_call_args assert ftp_mock.nlst.call_count == 2 assert ftp_mock.retrbinary.call_args_list[0][0][ 0] == 'RETR ' + model_file_path_full assert ftp_mock.retrbinary.call_args_list[1][0][ 0] == 'RETR ' + subfile_path_full
def test_download_artifacts_single(ftp_mock): repo = FTPArtifactRepository("ftp://test_ftp/some/path") repo.get_ftp_client = MagicMock() call_mock = MagicMock(return_value=ftp_mock) repo.get_ftp_client.return_value = MagicMock(__enter__=call_mock) ftp_mock.cwd = MagicMock(side_effect=ftplib.error_perm) repo.download_artifacts("test.txt") ftp_mock.retrbinary.assert_called_once() assert ftp_mock.retrbinary.call_args_list[0][0][0] == 'RETR /some/path/test.txt'