def test_list_artifacts(ftp_mock): artifact_root_path = "/experiment_id/run_id/" 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 # |- file # |- model # |- model.pb file_path = "file" file_size = 678 dir_path = "model" ftp_mock.cwd = MagicMock(side_effect=[None, ftplib.error_perm, None]) ftp_mock.nlst = MagicMock(return_value=[file_path, dir_path]) ftp_mock.size = MagicMock(return_value=file_size) artifacts = repo.list_artifacts(path=None) ftp_mock.nlst.assert_called_once_with(artifact_root_path) ftp_mock.size.assert_called_once_with(artifact_root_path + file_path) assert len(artifacts) == 2 assert artifacts[0].path == file_path assert artifacts[0].is_dir is False assert artifacts[0].file_size == file_size assert artifacts[1].path == dir_path assert artifacts[1].is_dir is True assert artifacts[1].file_size is None
def test_list_artifacts_empty(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.nlst = MagicMock(return_value=[]) assert repo.list_artifacts() == [] ftp_mock.nlst.assert_called_once_with("/some/path")
def test_list_artifacts_with_subdir(ftp_mock): artifact_root_path = "/experiment_id/run_id/" repo = FTPArtifactRepository("sftp://test_sftp" + 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 dir_name = 'model' # list artifacts at sub directory level file_path = 'model.pb' file_size = 345 subdir_name = 'variables' ftp_mock.nlst = MagicMock(return_value=[file_path, subdir_name]) ftp_mock.cwd = MagicMock(side_effect=[None, ftplib.error_perm, None]) ftp_mock.size = MagicMock(return_value=file_size) artifacts = repo.list_artifacts(path=dir_name) ftp_mock.nlst.assert_called_once_with(artifact_root_path + dir_name) ftp_mock.size.assert_called_once_with(artifact_root_path + dir_name + '/' + file_path) assert len(artifacts) == 2 assert artifacts[0].path == dir_name + '/' + file_path assert artifacts[0].is_dir is False assert artifacts[0].file_size == file_size assert artifacts[1].path == dir_name + '/' + subdir_name assert artifacts[1].is_dir is True assert artifacts[1].file_size is None