def test_list_artifacts(sftp_mock):
artifact_root_path = "/experiment_id/run_id/"
repo = SFTPArtifactRepository("sftp://test_sftp"+artifact_root_path, sftp_mock)
# mocked file structure
# |- file
# |- model
# |- model.pb
file_path = "file"
file_size = 678
dir_path = "model"
sftp_mock.isdir = MagicMock(side_effect=lambda path: {
(artifact_root_path+file_path): False,
(artifact_root_path+dir_path): True
}[path])
sftp_mock.listdir = MagicMock(return_value=[file_path, dir_path])
file_stat = MagicMock()
file_stat.configure_mock(st_size=file_size)
sftp_mock.stat = MagicMock(return_value=file_stat)
artifacts = repo.list_artifacts(path=None)
sftp_mock.listdir.assert_called_once_with(artifact_root_path)
sftp_mock.stat.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_with_subdir(sftp_mock):
artifact_root_path = "/experiment_id/run_id/"
repo = SFTPArtifactRepository("sftp://test_sftp" + artifact_root_path,
sftp_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'
sftp_mock.listdir = MagicMock(return_value=[file_path, subdir_name])
sftp_mock.isdir = MagicMock(
side_effect=lambda path: {
posixpath.join(artifact_root_path, dir_name): True,
posixpath.join(artifact_root_path, dir_name, file_path): False,
posixpath.join(artifact_root_path, dir_name, subdir_name): True,
}[path])
file_stat = MagicMock()
file_stat.configure_mock(st_size=file_size)
sftp_mock.stat = MagicMock(return_value=file_stat)
artifacts = repo.list_artifacts(path=dir_name)
sftp_mock.listdir.assert_called_once_with(artifact_root_path + dir_name)
sftp_mock.stat.assert_called_once_with(artifact_root_path + dir_name +
'/' + file_path)
assert len(artifacts) == 2
assert artifacts[0].path == posixpath.join(dir_name, file_path)
assert artifacts[0].is_dir is False
assert artifacts[0].file_size == file_size
assert artifacts[1].path == posixpath.join(dir_name, subdir_name)
assert artifacts[1].is_dir is True
assert artifacts[1].file_size is None
def test_list_artifacts_empty(sftp_mock):
repo = SFTPArtifactRepository("sftp://test_sftp/some/path", sftp_mock)
sftp_mock.listdir = MagicMock(return_value=[])
assert repo.list_artifacts() == []
sftp_mock.listdir.assert_called_once_with("/some/path")
