def test_download_artifact_throws_value_error_when_listed_blobs_do_not_contain_artifact_root_prefix( mock_client): repo = AzureBlobArtifactRepository(TEST_URI, mock_client) # Create a "bad blob" with a name that is not prefixed by the root path of the artifact store bad_blob_props = BlobProperties() bad_blob_props.size = 42 bad_blob_props.name = "file_path" def get_mock_listing(*args, **kwargs): """ Produces a mock listing that only contains content if the specified prefix is the artifact root. This allows us to mock `list_artifacts` during the `_download_artifacts_into` subroutine without recursively listing the same artifacts at every level of the directory traversal. """ # pylint: disable=unused-argument if posixpath.abspath(kwargs["name_starts_with"]) == posixpath.abspath( TEST_ROOT_PATH): # Return a blob that is not prefixed by the root path of the artifact store. This # should result in an exception being raised return MockBlobList([bad_blob_props]) else: return MockBlobList([]) mock_client.get_container_client( ).walk_blobs.side_effect = get_mock_listing with pytest.raises(MlflowException) as exc: repo.download_artifacts("") assert "Azure blob does not begin with the specified artifact path" in str( exc)
def test_log_artifacts(mock_client, tmpdir): repo = AzureBlobArtifactRepository(TEST_URI, mock_client) parentd = tmpdir.mkdir("data") subd = parentd.mkdir("subdir") parentd.join("a.txt").write("A") subd.join("b.txt").write("B") subd.join("c.txt").write("C") repo.log_artifacts(parentd.strpath) mock_client.get_container_client.assert_called_with("container") call_list = mock_client.get_container_client().upload_blob.call_args_list # Ensure that the order of the calls do not matter for call in call_list: arg1, arg2 = call[0] assert arg1 in [ posixpath.join(TEST_ROOT_PATH, x) for x in ["a.txt", "subdir/b.txt", "subdir/c.txt"] ] # arg2 should be a filebuffer if arg1.endswith("/a.txt"): assert arg2.name == os.path.normpath(parentd.strpath + "/a.txt") elif arg1.endswith("/b.txt"): assert arg2.name == os.path.normpath(subd.strpath + "/b.txt") elif arg1.endswith("/c.txt"): assert arg2.name == os.path.normpath(subd.strpath + "/c.txt") else: # This should be unreachable assert False
def test_list_artifacts_single_file(mock_client): repo = AzureBlobArtifactRepository(TEST_URI, mock_client) # Evaluate single file blob_props = BlobProperties() blob_props.name = posixpath.join(TEST_ROOT_PATH, "file") mock_client.get_container_client().walk_blobs.return_value = MockBlobList( [blob_props]) assert repo.list_artifacts("file") == []
def test_download_directory_artifact_succeeds_when_artifact_root_is_blob_container_root( mock_client, tmpdir): repo = AzureBlobArtifactRepository(TEST_BLOB_CONTAINER_ROOT, mock_client) subdir_path = "my_directory" dir_prefix = BlobPrefix() dir_prefix.name = subdir_path file_path_1 = "file_1" file_path_2 = "file_2" blob_props_1 = BlobProperties() blob_props_1.size = 42 blob_props_1.name = posixpath.join(subdir_path, file_path_1) blob_props_2 = BlobProperties() blob_props_2.size = 42 blob_props_2.name = posixpath.join(subdir_path, file_path_2) def get_mock_listing(*args, **kwargs): """ Produces a mock listing that only contains content if the specified prefix is the artifact root or a relevant subdirectory. This allows us to mock `list_artifacts` during the `_download_artifacts_into` subroutine without recursively listing the same artifacts at every level of the directory traversal. """ # pylint: disable=unused-argument if posixpath.abspath(kwargs["name_starts_with"]) == "/": return MockBlobList([dir_prefix]) if posixpath.abspath( kwargs["name_starts_with"]) == posixpath.abspath(subdir_path): return MockBlobList([blob_props_1, blob_props_2]) else: return MockBlobList([]) def create_file(buffer): fname = os.path.basename(buffer.name) f = tmpdir.join(fname) f.write("hello world!") mock_client.get_container_client( ).walk_blobs.side_effect = get_mock_listing mock_client.get_container_client().download_blob( ).readinto.side_effect = create_file # Ensure that the root directory can be downloaded successfully repo.download_artifacts("") # Ensure that the `mkfile` side effect copied all of the download artifacts into `tmpdir` dir_contents = os.listdir(tmpdir.strpath) assert file_path_1 in dir_contents assert file_path_2 in dir_contents
def test_log_artifact(mock_client, tmpdir): repo = AzureBlobArtifactRepository(TEST_URI, mock_client) d = tmpdir.mkdir("data") f = d.join("test.txt") f.write("hello world!") fpath = posixpath.join(d.strpath, "test.txt") repo.log_artifact(fpath) mock_client.get_container_client.assert_called_with("container") arg1, arg2 = mock_client.get_container_client().upload_blob.call_args[0] assert arg1 == posixpath.join(TEST_ROOT_PATH, "test.txt") # arg2 should be a filebuffer assert arg2.name == fpath
def test_download_file_artifact(mock_client, tmpdir): repo = AzureBlobArtifactRepository(TEST_URI, mock_client) mock_client.get_container_client().walk_blobs.return_value = MockBlobList( []) def create_file(buffer): local_path = os.path.basename(buffer.name) f = tmpdir.join(local_path) f.write("hello world!") mock_client.get_container_client().download_blob( ).readinto.side_effect = create_file repo.download_artifacts("test.txt") assert os.path.exists(os.path.join(tmpdir.strpath, "test.txt")) mock_client.get_container_client().download_blob.assert_called_with( posixpath.join(TEST_ROOT_PATH, "test.txt"))
def test_list_artifacts(mock_client): repo = AzureBlobArtifactRepository(TEST_URI, mock_client) # Create some files to return dir_prefix = BlobPrefix() dir_prefix.name = posixpath.join(TEST_ROOT_PATH, "dir") blob_props = BlobProperties() blob_props.size = 42 blob_props.name = posixpath.join(TEST_ROOT_PATH, "file") mock_client.get_container_client().walk_blobs.return_value = MockBlobList( [dir_prefix, blob_props]) artifacts = repo.list_artifacts() assert artifacts[0].path == "dir" assert artifacts[0].is_dir is True assert artifacts[0].file_size is None assert artifacts[1].path == "file" assert artifacts[1].is_dir is False assert artifacts[1].file_size == 42
def test_list_artifacts_empty(mock_client): repo = AzureBlobArtifactRepository(TEST_URI, mock_client) mock_client.get_container_client().walk_blobs.return_value = MockBlobList( []) assert repo.list_artifacts() == []
def test_exception_if_no_env_vars(mock_client): # pylint: disable=unused-argument # We pass in the mock_client here to clear Azure environment variables, but we don't use it with pytest.raises(Exception, match="AZURE_STORAGE_CONNECTION_STRING"): AzureBlobArtifactRepository(TEST_URI)