diff --git a/tests/func/data/db/test_index.py b/tests/func/data/db/test_index.py index d1bfab69c..eb077fb89 100644 --- a/tests/func/data/db/test_index.py +++ b/tests/func/data/db/test_index.py @@ -1,7 +1,8 @@ +import os + import pytest from dvc.exceptions import DownloadError, UploadError -from dvc.fs import LocalFileSystem from dvc.utils.fs import remove from dvc_data.hashfile.db import get_index @@ -63,22 +64,54 @@ def test_clear_on_download_err(tmp_dir, dvc, index, mocker): assert list(index.hashes()) - mocker.patch("dvc_objects.fs.generic.transfer", side_effect=Exception) + def unreliable_download(_from_fs, from_info, _to_fs, to_info, **kwargs): + on_error = kwargs["on_error"] + assert on_error + if isinstance(from_info, str): + from_info = [from_info] + if isinstance(to_info, str): + to_info = [to_info] + for from_i, to_i in zip(from_info, to_info): + on_error(from_i, to_i, Exception()) + + mocker.patch("dvc_objects.fs.generic.transfer", unreliable_download) with pytest.raises(DownloadError): dvc.pull() assert not list(index.hashes()) def test_partial_upload(tmp_dir, dvc, index, mocker): + from dvc_objects.fs import generic + tmp_dir.dvc_gen({"foo": "foo content"}) - tmp_dir.dvc_gen({"bar": {"baz": "baz content"}}) + baz = tmp_dir.dvc_gen({"bar": {"baz": "baz content"}})[0].outs[0] - original = LocalFileSystem.put_file + original = generic.transfer + odb = dvc.cloud.get_remote_odb("upstream") - def unreliable_upload(self, from_file, to_info, name=None, **kwargs): - if "baz" in name: - raise Exception("stop baz") - return original(self, from_file, to_info, name, **kwargs) + def unreliable_upload(from_fs, from_info, to_fs, to_info, **kwargs): + on_error = kwargs["on_error"] + assert on_error + if isinstance(from_info, str): + from_info = [from_info] + else: + from_info = list(from_info) + if isinstance(to_info, str): + to_info = [to_info] + else: + to_info = list(to_info) + for i in range(len(from_info) - 1, -1, -1): + from_i = from_info[i] + to_i = to_info[i] + if os.path.abspath(to_i) == os.path.abspath( + odb.get(baz.hash_info.value).path + ): + if on_error: + on_error(from_i, to_i, Exception("stop baz")) + del from_info[i] + del to_info[i] + + return original(from_fs, from_info, to_fs, to_info, **kwargs) mocker.patch("dvc_objects.fs.generic.transfer", unreliable_upload) with pytest.raises(UploadError): diff --git a/tests/func/test_import.py b/tests/func/test_import.py index 1ac692173..ec35cf97e 100644 --- a/tests/func/test_import.py +++ b/tests/func/test_import.py @@ -369,7 +369,17 @@ def test_download_error_pulling_imported_stage( remove("foo_imported") remove(dst_cache) - mocker.patch("dvc_objects.fs.generic.transfer", side_effect=Exception) + def unreliable_download(_from_fs, from_info, _to_fs, to_info, **kwargs): + on_error = kwargs["on_error"] + assert on_error + if isinstance(from_info, str): + from_info = [from_info] + if isinstance(to_info, str): + to_info = [to_info] + for from_i, to_i in zip(from_info, to_info): + on_error(from_i, to_i, Exception()) + + mocker.patch("dvc_objects.fs.generic.transfer", unreliable_download) with pytest.raises(DownloadError): dvc.pull(["foo_imported.dvc"]) diff --git a/tests/func/test_remote.py b/tests/func/test_remote.py index dce848a26..7f3b1af1d 100644 --- a/tests/func/test_remote.py +++ b/tests/func/test_remote.py @@ -1,4 +1,5 @@ import errno +import itertools import os import stat from unittest.mock import patch @@ -163,7 +164,9 @@ def test_dir_hash_should_be_key_order_agnostic(tmp_dir, dvc): assert hash1 == hash2 -def test_partial_push_n_pull(tmp_dir, dvc, tmp_path_factory, local_remote): +def test_partial_push_n_pull( # noqa: C901 + tmp_dir, dvc, tmp_path_factory, local_remote +): from dvc_objects.fs import generic foo = tmp_dir.dvc_gen({"foo": "foo content"})[0].outs[0] @@ -175,10 +178,26 @@ def test_partial_push_n_pull(tmp_dir, dvc, tmp_path_factory, local_remote): odb = dvc.cloud.get_remote_odb("upstream") def unreliable_upload(from_fs, from_info, to_fs, to_info, **kwargs): - if os.path.abspath(to_info) == os.path.abspath( - odb.get(foo.hash_info.value).path - ): - raise Exception("stop foo") + on_error = kwargs["on_error"] + assert on_error + if isinstance(from_info, str): + from_info = [from_info] + else: + from_info = list(from_info) + if isinstance(to_info, str): + to_info = [to_info] + else: + to_info = list(to_info) + for i in range(len(from_info) - 1, -1, -1): + from_i = from_info[i] + to_i = to_info[i] + if os.path.abspath(to_i) == os.path.abspath( + odb.get(foo.hash_info.value).path + ): + if on_error: + on_error(from_i, to_i, Exception("stop foo")) + del from_info[i] + del to_info[i] return original(from_fs, from_info, to_fs, to_info, **kwargs) with patch.object(generic, "transfer", unreliable_upload): @@ -195,7 +214,18 @@ def test_partial_push_n_pull(tmp_dir, dvc, tmp_path_factory, local_remote): dvc.odb.local.clear() baz._collect_used_dir_cache() - with patch.object(generic, "transfer", side_effect=Exception): + + def unreliable_download(_from_fs, from_info, _to_fs, to_info, **kwargs): + on_error = kwargs["on_error"] + assert on_error + if isinstance(from_info, str): + from_info = [from_info] + if isinstance(to_info, str): + to_info = [to_info] + for from_i, to_i in zip(from_info, to_info): + on_error(from_i, to_i, Exception()) + + with patch.object(generic, "transfer", unreliable_download): with pytest.raises(DownloadError) as download_error_info: dvc.pull() # error count should be len(.dir + standalone file checksums) @@ -270,7 +300,11 @@ def test_push_order(tmp_dir, dvc, tmp_path_factory, mocker, local_remote): odb = dvc.cloud.get_remote_odb("upstream") foo_path = odb.oid_to_path(foo.hash_info.value) bar_path = odb.oid_to_path(foo.obj._trie[("bar",)][1].value) - paths = [args[3] for args, _ in mocked_upload.call_args_list] + paths = list( + itertools.chain.from_iterable( + args[3] for args, _ in mocked_upload.call_args_list + ) + ) assert paths.index(foo_path) > paths.index(bar_path)