From 6cb37ac27b6224cb4913de9801ac58075e969645 Mon Sep 17 00:00:00 2001 From: skshetry <18718008+skshetry@users.noreply.github.com> Date: Mon, 26 Feb 2024 21:45:09 +0545 Subject: [PATCH] tests: add tests for datasets (#10321) --- dvc/commands/dataset.py | 20 +- dvc/repo/__init__.py | 1 + dvc/repo/datasets.py | 82 ++++-- dvc/repo/index.py | 5 - tests/func/test_dataset.py | 414 +++++++++++++++++++++++++++++ tests/unit/command/test_dataset.py | 127 +++++++++ tests/unit/test_dataset.py | 45 ++++ 7 files changed, 651 insertions(+), 43 deletions(-) create mode 100644 tests/func/test_dataset.py create mode 100644 tests/unit/command/test_dataset.py create mode 100644 tests/unit/test_dataset.py diff --git a/dvc/commands/dataset.py b/dvc/commands/dataset.py index 9cbc95c33..f92238e49 100644 --- a/dvc/commands/dataset.py +++ b/dvc/commands/dataset.py @@ -49,22 +49,6 @@ class CmdDatasetAdd(CmdBase): ui.write(action, ui.rich_text(name, "cyan"), text, styled=True) def run(self): - from urllib.parse import urlsplit - - d = vars(self.args) - url_obj = urlsplit(self.args.url) - if url_obj.scheme == "dvcx": - d["type"] = "dvcx" - elif url_obj.scheme.startswith("dvc"): - d["type"] = "dvc" - protos = tuple(url_obj.scheme.split("+")) - if not protos or protos == ("dvc",) or protos == ("dvc", "ssh"): - d["url"] = url_obj.netloc + url_obj.path - else: - d["url"] = url_obj._replace(scheme=protos[1]).geturl() - else: - d["type"] = "url" - existing = self.repo.datasets.get(self.args.name) with self.repo.scm_context: if not self.args.force and existing: @@ -73,8 +57,9 @@ class CmdDatasetAdd(CmdBase): f"{self.args.name} already exists in {path}, " "use the --force to overwrite" ) - dataset = self.repo.datasets.add(**d) + dataset = self.repo.datasets.add(**vars(self.args)) self.display(self.args.name, dataset) + return 0 class CmdDatasetUpdate(CmdBase): @@ -133,6 +118,7 @@ class CmdDatasetUpdate(CmdBase): ) return 1 self.display(self.args.name, dataset, new) + return 0 def add_parser(subparsers, parent_parser): diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py index 3e14069d4..7ca79a74d 100644 --- a/dvc/repo/__init__.py +++ b/dvc/repo/__init__.py @@ -655,6 +655,7 @@ class Repo: def _reset(self): self.scm._reset() + self.datasets._reset() self.state.close() if "dvcfs" in self.__dict__: self.dvcfs.close() diff --git a/dvc/repo/datasets.py b/dvc/repo/datasets.py index 4188d9f5a..1131d5d94 100644 --- a/dvc/repo/datasets.py +++ b/dvc/repo/datasets.py @@ -1,3 +1,4 @@ +import os from collections.abc import Iterator, Mapping from datetime import datetime from functools import cached_property @@ -11,10 +12,12 @@ from attrs.converters import default_if_none from dvc.dvcfile import Lockfile, ProjectFile from dvc.exceptions import DvcException from dvc.log import logger +from dvc.types import StrPath from dvc_data.hashfile.meta import Meta if TYPE_CHECKING: from dql.dataset import DatasetRecord # type: ignore[import] + from dvcx.dataset import DatasetVersion # type: ignore[import] from typing_extensions import Self from dvc.repo import Repo @@ -23,6 +26,50 @@ if TYPE_CHECKING: logger = logger.getChild(__name__) +def parse_url_and_type(url: str): + from urllib.parse import urlsplit + + if os.path.exists(url): + return {"type": "dvc", "url": url} + + url_obj = urlsplit(url) + if url_obj.scheme == "dvcx": + return {"type": "dvcx", "url": url} + if url_obj.scheme and not url_obj.scheme.startswith("dvc"): + return {"type": "url", "url": url} + + protos = tuple(url_obj.scheme.split("+")) + if not protos or protos == ("dvc",): + url = url_obj.netloc + url_obj.path + else: + url = url_obj._replace(scheme=protos[1]).geturl() + return {"type": "dvc", "url": url} + + +def _get_dataset_record(name: str) -> "DatasetRecord": + from dvc.exceptions import DvcException + + try: + from dvcx.catalog import get_catalog # type: ignore[import] + + except ImportError as exc: + raise DvcException("dvcx is not installed") from exc + + catalog = get_catalog() + return catalog.get_remote_dataset(name) + + +def _get_dataset_info( + name: str, record: Optional["DatasetRecord"] = None, version: Optional[int] = None +) -> "DatasetVersion": + record = record or _get_dataset_record(name) + assert record + v = version or record.latest_version + assert v + assert v >= 1 + return record.get_version(v) + + def default_str(v) -> str: return default_if_none("")(v) @@ -160,22 +207,9 @@ class DVCXDataset: version: Optional[int] = None, **kwargs, ) -> "Self": - if not record: - try: - from dvcx.catalog import get_catalog # type: ignore[import] - - except ImportError as exc: - raise DvcException("dvcx is not installed") from exc - - name, _version = self.name_version - version = _version or version - catalog = get_catalog() - record = catalog.get_remote_dataset(name) - - assert record is not None - ver = version or record.latest_version - assert ver - version_info = record.get_version(ver) + name, _version = self.name_version + version = version or _version + version_info = _get_dataset_info(name, record=record, version=version) lock = DVCXDatasetLock( **self.spec.to_dict(), version=version_info.version, @@ -264,6 +298,8 @@ class Datasets(Mapping[str, Dataset]): datasets_lock = self.repo.index._datasets_lock def find(path, name) -> Optional[dict[str, Any]]: + # only look for `name` in the lock file next to the + # corresponding `dvc.yaml` file lock = datasets_lock.get(path, []) return next((dataset for dataset in lock if dataset["name"] == name), None) @@ -348,10 +384,14 @@ class Datasets(Mapping[str, Dataset]): raise ValueError(f"unknown dataset type: {spec.type!r}") def add( - self, url: str, name: str, manifest_path: str = "dvc.yaml", **kwargs: Any + self, + url: str, + name: str, + manifest_path: StrPath = "dvc.yaml", + **kwargs: Any, ) -> Dataset: - kwargs.update({"url": url, "name": name}) - dataset = self._build_dataset(manifest_path, kwargs) + spec = kwargs | parse_url_and_type(url) | {"name": name} + dataset = self._build_dataset(os.path.abspath(manifest_path), spec) dataset = dataset.update(self.repo) self.dump(dataset) @@ -366,13 +406,13 @@ class Datasets(Mapping[str, Dataset]): self[name] = new return dataset, new - def _dump_spec(self, manifest_path: str, spec: Spec) -> None: + def _dump_spec(self, manifest_path: StrPath, spec: Spec) -> None: spec_data = spec.to_dict() assert spec_data.keys() & {"type", "name", "url"} project_file = ProjectFile(self.repo, manifest_path) project_file.dump_dataset(spec_data) - def _dump_lock(self, manifest_path: str, lock: Lock) -> None: + def _dump_lock(self, manifest_path: StrPath, lock: Lock) -> None: lock_data = lock.to_dict() assert lock_data.keys() & {"type", "name", "url"} lockfile = Lockfile(self.repo, Path(manifest_path).with_suffix(".lock")) diff --git a/dvc/repo/index.py b/dvc/repo/index.py index 2f764a786..a3af9736b 100644 --- a/dvc/repo/index.py +++ b/dvc/repo/index.py @@ -414,11 +414,6 @@ class Index: for stage in self.stages: yield from stage.outs - @cached_property - def datasets(self) -> dict[str, dict[str, Any]]: - datasets = chain.from_iterable(self._datasets.values()) - return {dataset["name"]: dataset for dataset in datasets} - @cached_property def out_data_keys(self) -> dict[str, set["DataIndexKey"]]: by_workspace: dict[str, set["DataIndexKey"]] = defaultdict(set) diff --git a/tests/func/test_dataset.py b/tests/func/test_dataset.py new file mode 100644 index 000000000..ad1f9be85 --- /dev/null +++ b/tests/func/test_dataset.py @@ -0,0 +1,414 @@ +import os +from datetime import datetime, timezone +from typing import TYPE_CHECKING + +import pytest +from attrs import define, evolve, field, has + +from dvc.dependency.base import Dependency +from dvc.exceptions import ReproductionError +from dvc.repo.datasets import ( + DatasetNotFoundError, + DatasetSpec, + DVCDataset, + DVCDatasetLock, + DVCDatasetSpec, + DVCXDataset, + DVCXDatasetLock, + FileInfo, + URLDataset, + URLDatasetLock, +) +from dvc_data.hashfile.meta import Meta +from dvc_data.index import HashInfo, Tree + +if TYPE_CHECKING: + from dvc.repo import Repo + + +@define +class MockedDVCXVersionInfo: + version: int + created_at: datetime = field(factory=lambda: datetime.now(timezone.utc)) + + +def evolve_recursive(inst, **changes): + """Recursive attr.evolve() method, where any attr-based attributes + will be evolved too. + """ + for key, value in changes.items(): + v = getattr(inst, key) + if has(type(v)) and isinstance(value, dict): + value = evolve_recursive(v, **value) + changes[key] = value + return evolve(inst, **changes) + + +def test_dvc(tmp_dir, scm, dvc: "Repo"): + datasets = dvc.datasets + + tmp_dir.scm_gen("file", "file", commit="add file") + dataset = datasets.add(tmp_dir.fs_path, name="mydataset", path="file") + expected = DVCDataset( + manifest_path=(tmp_dir / "dvc.yaml").fs_path, + spec=DVCDatasetSpec( + name="mydataset", url=tmp_dir.fs_path, type="dvc", path="file" + ), + lock=DVCDatasetLock( + name="mydataset", + url=tmp_dir.fs_path, + type="dvc", + path="file", + rev_lock=scm.get_rev(), + ), + ) + assert "mydataset" in datasets + assert dataset == datasets["mydataset"] == expected + tmp_dir.scm_gen("file", "file", commit="update file") + + old, new = datasets.update("mydataset") + assert old == dataset + assert old != new + expected = evolve_recursive(expected, lock={"rev_lock": scm.get_rev()}) + assert new == datasets["mydataset"] == expected + + # noop + old, new = datasets.update("mydataset") + assert old == new + + +def test_dvcx(tmp_dir, dvc, mocker): + datasets = dvc.datasets + + version_info = [MockedDVCXVersionInfo(1), MockedDVCXVersionInfo(2)] + version_info.append(version_info[1]) + mocker.patch("dvc.repo.datasets._get_dataset_info", side_effect=version_info) + + dataset = datasets.add(url="dvcx://dataset", name="mydataset") + expected = DVCXDataset( + manifest_path=(tmp_dir / "dvc.yaml").fs_path, + spec=DatasetSpec(name="mydataset", url="dvcx://dataset", type="dvcx"), + lock=DVCXDatasetLock( + name="mydataset", + url="dvcx://dataset", + type="dvcx", + version=1, + created_at=version_info[0].created_at, + ), + ) + assert "mydataset" in datasets + assert dataset == datasets["mydataset"] == expected + + old, new = datasets.update("mydataset") + assert old == dataset + assert old != new + expected = evolve_recursive( + expected, lock={"version": 2, "created_at": version_info[1].created_at} + ) + assert new == datasets["mydataset"] == expected + + # noop + old, new = datasets.update("mydataset") + assert old == new + + +def test_url(tmp_dir, dvc, mocker): + datasets = dvc.datasets + + tree = Tree() + tree.add(("foo",), Meta(version_id="1"), None) + tree.digest(with_meta=True) + tree_meta = Meta(isdir=True) + + def mocked_save(d): + d.meta, d.obj, d.hash_info = tree_meta, tree, HashInfo("md5", "value.dir") + + mocker.patch.object(Dependency, "save", mocked_save) + + dataset = datasets.add(url="s3://dataset", name="mydataset") + expected = URLDataset( + manifest_path=(tmp_dir / "dvc.yaml").fs_path, + spec=DatasetSpec(name="mydataset", url="s3://dataset", type="url"), + lock=URLDatasetLock( + name="mydataset", + url="s3://dataset", + type="url", + meta=Meta(isdir=True), + files=[FileInfo(relpath="foo", meta=Meta(version_id="1"))], + ), + ) + assert "mydataset" in datasets + assert dataset == datasets["mydataset"] == expected + + tree.add(("bar",), Meta(version_id="2"), None) + old, new = datasets.update("mydataset") + assert old == dataset + assert old != new + + assert expected.lock + new_files = [ + *expected.lock.files, + FileInfo(relpath="bar", meta=Meta(version_id="2")), + ] + + expected = evolve_recursive(expected, lock={"files": new_files}) + assert new == datasets["mydataset"] == expected + + # noop + old, new = datasets.update("mydataset") + assert old == new + + +def test_dvc_dump(tmp_dir, dvc): + manifest_path = os.path.join(tmp_dir, "dvc.yaml") + spec = DVCDatasetSpec( + name="mydataset", url=tmp_dir.fs_path, type="dvc", path="path", rev="main" + ) + lock = DVCDatasetLock(rev_lock="0" * 40, **spec.to_dict()) + dataset = DVCDataset(manifest_path=manifest_path, spec=spec, lock=lock) + + dvc.datasets.dump(dataset) + + spec_d = { + "name": "mydataset", + "type": "dvc", + "url": tmp_dir.fs_path, + "path": "path", + "rev": "main", + } + assert (tmp_dir / "dvc.yaml").parse() == {"datasets": [spec_d]} + assert (tmp_dir / "dvc.lock").parse() == { + "schema": "2.0", + "stages": {}, + "datasets": [{**spec_d, "rev_lock": "0" * 40}], + } + + dvc._reset() + assert "_datasets" not in vars(dvc.datasets) + # test that we can read them back + assert dvc.datasets["mydataset"] == dataset + + +def test_dvcx_dump(tmp_dir, dvc): + manifest_path = os.path.join(tmp_dir, "dvc.yaml") + spec = DatasetSpec(name="mydataset", url="dvcx://dataset", type="dvcx") + dt = datetime.now(tz=timezone.utc) + lock = DVCXDatasetLock(version=1, created_at=dt, **spec.to_dict()) + dataset = DVCXDataset(manifest_path=manifest_path, spec=spec, lock=lock) + + dvc.datasets.dump(dataset) + + spec_d = {"name": "mydataset", "type": "dvcx", "url": "dvcx://dataset"} + assert (tmp_dir / "dvc.yaml").parse() == {"datasets": [spec_d]} + assert (tmp_dir / "dvc.lock").parse() == { + "schema": "2.0", + "stages": {}, + "datasets": [{**spec_d, "version": 1, "created_at": dt.isoformat()}], + } + + dvc._reset() + assert "_datasets" not in vars(dvc.datasets) + # test that we can read them back + assert dvc.datasets["mydataset"] == dataset + + +def test_url_dump(tmp_dir, dvc): + manifest_path = os.path.join(tmp_dir, "dvc.yaml") + spec = DatasetSpec(name="mydataset", url="s3://dataset", type="url") + files = [FileInfo(relpath="foo", meta=Meta(version_id="1"))] + lock = URLDatasetLock(meta=Meta(isdir=True), files=files, **spec.to_dict()) + dataset = URLDataset(manifest_path=manifest_path, spec=spec, lock=lock) + + dvc.datasets.dump(dataset) + + spec_d = {"name": "mydataset", "url": "s3://dataset", "type": "url"} + assert (tmp_dir / "dvc.yaml").parse() == {"datasets": [spec_d]} + assert (tmp_dir / "dvc.lock").parse() == { + "schema": "2.0", + "datasets": [ + { + **spec_d, + "meta": {"isdir": True}, + "files": [{"relpath": "foo", "meta": {"version_id": "1"}}], + } + ], + "stages": {}, + } + + dvc._reset() + assert "_datasets" not in vars(dvc.datasets) + # test that we can read them back + assert dvc.datasets["mydataset"] == dataset + + +def test_invalidation(tmp_dir, dvc): + manifest_path = os.path.join(tmp_dir, "dvc.yaml") + spec = DatasetSpec(name="mydataset", url="url1", type="url") + lock = DVCXDatasetLock( + name="mydataset", + url="dvcx://dataset", + type="dvcx", + version=1, + created_at=datetime.now(tz=timezone.utc), + ) + dvc.datasets._dump_spec(manifest_path, spec) + dvc.datasets._dump_lock(manifest_path, lock) + + assert dvc.datasets["mydataset"] == URLDataset( + manifest_path=manifest_path, + spec=spec, + lock=None, # lock should be discarded + ) + + +def test_dvc_dataset_pipeline(tmp_dir, dvc, scm): + dvc.datasets.add(name="mydataset", url=tmp_dir.fs_path) + + stage = dvc.stage.add(cmd="echo", name="train", deps=["ds://mydataset"]) + assert (tmp_dir / "dvc.yaml").parse() == { + "datasets": [{"name": "mydataset", "url": tmp_dir.fs_path, "type": "dvc"}], + "stages": {"train": {"cmd": "echo", "deps": ["ds://mydataset"]}}, + } + + assert dvc.status() == {"train": [{"changed deps": {"ds://mydataset": "new"}}]} + assert dvc.reproduce() == [stage] + + d = (tmp_dir / "dvc.lock").parse() + assert d["stages"]["train"]["deps"][0] == { + "path": "ds://mydataset", + "dataset": d["datasets"][0], + } + + assert dvc.status() == {} + assert dvc.reproduce() == [] + + tmp_dir.scm_gen("foo", "foo", commit="add foo") + dvc.datasets.update("mydataset") + + assert dvc.status() == {"train": [{"changed deps": {"ds://mydataset": "modified"}}]} + assert dvc.reproduce() == [stage] + + +def test_dvcx_dataset_pipeline(mocker, tmp_dir, dvc): + version_info = [MockedDVCXVersionInfo(1), MockedDVCXVersionInfo(2)] + mocker.patch("dvc.repo.datasets._get_dataset_info", side_effect=version_info) + + dvc.datasets.add(name="mydataset", url="dvcx://mydataset") + + stage = dvc.stage.add(cmd="echo", name="train", deps=["ds://mydataset"]) + assert (tmp_dir / "dvc.yaml").parse() == { + "datasets": [{"name": "mydataset", "url": "dvcx://mydataset", "type": "dvcx"}], + "stages": {"train": {"cmd": "echo", "deps": ["ds://mydataset"]}}, + } + + assert dvc.status() == {"train": [{"changed deps": {"ds://mydataset": "new"}}]} + assert dvc.reproduce() == [stage] + + d = (tmp_dir / "dvc.lock").parse() + assert d["stages"]["train"]["deps"][0] == { + "path": "ds://mydataset", + "dataset": d["datasets"][0], + } + + assert dvc.status() == {} + assert dvc.reproduce() == [] + + dvc.datasets.update("mydataset") + + assert dvc.status() == {"train": [{"changed deps": {"ds://mydataset": "modified"}}]} + assert dvc.reproduce() == [stage] + + +def test_url_dataset_pipeline(mocker, tmp_dir, dvc): + tree = Tree() + tree.add(("foo",), Meta(version_id="1"), None) + tree.digest(with_meta=True) + tree_meta = Meta(isdir=True) + + def mocked_save(d): + d.meta, d.obj, d.hash_info = tree_meta, tree, HashInfo("md5", "value.dir") + + mocker.patch.object(Dependency, "save", mocked_save) + + dvc.datasets.add(name="mydataset", url="s3://mydataset") + + stage = dvc.stage.add(cmd="echo", name="train", deps=["ds://mydataset"]) + assert (tmp_dir / "dvc.yaml").parse() == { + "datasets": [{"name": "mydataset", "url": "s3://mydataset", "type": "url"}], + "stages": {"train": {"cmd": "echo", "deps": ["ds://mydataset"]}}, + } + + assert dvc.status() == {"train": [{"changed deps": {"ds://mydataset": "new"}}]} + assert dvc.reproduce() == [stage] + + d = (tmp_dir / "dvc.lock").parse() + assert d["stages"]["train"]["deps"][0] == { + "path": "ds://mydataset", + "dataset": d["datasets"][0], + } + + assert dvc.status() == {} + assert dvc.reproduce() == [] + + tree.add(("bar",), Meta(version_id="2"), None) + dvc.datasets.update("mydataset") + + assert dvc.status() == {"train": [{"changed deps": {"ds://mydataset": "modified"}}]} + assert dvc.reproduce() == [stage] + + +def test_pipeline_when_not_in_sync(tmp_dir, dvc): + manifest_path = os.path.join(tmp_dir, "dvc.yaml") + spec = DatasetSpec(name="mydataset", url="url1", type="url") + lock = DVCXDatasetLock( + name="mydataset", + url="dvcx://dataset", + type="dvcx", + version=1, + created_at=datetime.now(tz=timezone.utc), + ) + dvc.datasets._dump_spec(manifest_path, spec) + dvc.datasets._dump_lock(manifest_path, lock) + + dvc.stage.add(name="train", cmd="echo", deps=["ds://mydataset"]) + assert dvc.status() == { + "train": [{"changed deps": {"ds://mydataset": "not in sync"}}] + } + with pytest.raises(ReproductionError) as exc: + dvc.reproduce() + assert "not in sync" in str(exc.value.__cause__) + + +def test_collect(tmp_dir, dvc): + manifest_path1 = os.path.join(tmp_dir, "dvc.yaml") + dt = datetime.now(tz=timezone.utc) + spec = DatasetSpec(name="mydataset1", url="url1", type="dvcx") + lock = DVCXDatasetLock(version=1, created_at=dt, **spec.to_dict()) + mydataset1 = DVCXDataset(manifest_path=manifest_path1, spec=spec, lock=lock) + dvc.datasets.dump(mydataset1) + + (tmp_dir / "sub").mkdir() + manifest_path2 = os.path.join(tmp_dir, "sub", "dvc.yaml") + spec = DVCDatasetSpec( + name="mydataset2", url=tmp_dir.fs_path, type="dvc", path="path" + ) + lock = DVCDatasetLock(rev_lock="0" * 40, **spec.to_dict()) + mydataset2 = DVCDataset(manifest_path=manifest_path2, spec=spec, lock=lock) + dvc.datasets.dump(mydataset2) + + dvc._reset() + assert "_datasets" not in vars(dvc.datasets) + + assert len(dvc.datasets) == 2 + assert "mydataset1" in dvc.datasets + assert "mydataset2" in dvc.datasets + assert list(iter(dvc.datasets)) == ["mydataset1", "mydataset2"] + assert dvc.datasets["mydataset1"] == mydataset1 + assert dvc.datasets["mydataset2"] == mydataset2 + assert dict(dvc.datasets.items()) == { + "mydataset1": mydataset1, + "mydataset2": mydataset2, + } + + with pytest.raises(DatasetNotFoundError, match=r"^dataset not found$"): + dvc.datasets["not-existing"] diff --git a/tests/unit/command/test_dataset.py b/tests/unit/command/test_dataset.py new file mode 100644 index 000000000..776b1c521 --- /dev/null +++ b/tests/unit/command/test_dataset.py @@ -0,0 +1,127 @@ +from datetime import datetime + +import pytest + +from dvc.cli import main + + +@pytest.mark.parametrize( + "spec,lock,expected_output", + [ + ( + {"name": "ds", "url": "url", "type": "dvc", "path": "path", "rev": "main"}, + {"rev_lock": "0" * 40}, + "Adding ds (url:/path @ main)\n", + ), + ( + {"name": "mydataset", "url": "dvcx://dataset", "type": "dvcx"}, + {"rev_lock": "0" * 40, "version": 1, "created_at": datetime.now()}, + "Adding mydataset (dvcx://dataset @ v1)\n", + ), + ( + {"name": "mydataset", "url": "s3://bucket/path", "type": "url"}, + { + "files": [{"relpath": "foo", "meta": {"version_id": 1}}], + "meta": {"isdir": True}, + }, + "Adding mydataset (s3://bucket/path)\n", + ), + ], +) +def test_add(dvc, capsys, mocker, spec, lock, expected_output): + dataset = dvc.datasets._build_dataset("dvc.yaml", spec, spec | lock) + + m = mocker.patch("dvc.repo.datasets.Datasets.add", return_value=dataset) + + assert main(["dataset", "add", "--url", spec["url"], "--name", spec["name"]]) == 0 + out, err = capsys.readouterr() + assert out == expected_output + assert not err + m.assert_called_once() + + +def test_add_already_exists(dvc, caplog, mocker): + spec = {"name": "ds", "url": "url", "type": "dvc"} + dataset = dvc.datasets._build_dataset("dvc.yaml", spec, None) + mocker.patch("dvc.repo.datasets.Datasets.get", return_value=dataset) + + assert main(["dataset", "add", "--url", "url", "--name", "ds"]) == 255 + assert "ds already exists in dvc.yaml, use the --force to overwrite" in caplog.text + + +@pytest.mark.parametrize("lock", ["missing", "unchanged", "updated"]) +@pytest.mark.parametrize( + "spec,old_lock,new_lock,expected_outputs", + [ + ( + {"name": "ds", "url": "url", "type": "dvc", "path": "path", "rev": "main"}, + {"rev_lock": "0" * 40}, + {"rev_lock": "1" * 40}, + { + "missing": "Updating ds (url:/path @ main)\n", + "unchanged": "Nothing to update\n", + "updated": "Updating ds (000000000 -> 111111111)\n", + }, + ), + ( + {"name": "mydataset", "url": "dvcx://dataset", "type": "dvcx"}, + {"rev_lock": "0" * 40, "version": 1, "created_at": datetime.now()}, + {"version": 2}, + { + "missing": "Updating mydataset (dvcx://dataset @ v2)\n", + "unchanged": "Nothing to update\n", + "updated": "Updating mydataset (v1 -> v2)\n", + }, + ), + ( + {"name": "mydataset", "url": "s3://bucket/path", "type": "url"}, + { + "files": [ + {"relpath": "bar", "meta": {"version_id": 2}}, + {"relpath": "baz", "meta": {"version_id": 3}}, + {"relpath": "foo", "meta": {"version_id": 1}}, + ], + "meta": {"isdir": True}, + }, + { + "files": [ + {"relpath": "bar", "meta": {"version_id": 2}}, # unchanged + {"relpath": "baz", "meta": {"version_id": 4}}, # modified + # `foo` deleted + {"relpath": "foobar", "meta": {"version_id": 5}}, # new + ], + "meta": {"isdir": True}, + }, + { + "missing": "Updating mydataset (s3://bucket/path)\n", + "unchanged": "Nothing to update\n", + "updated": ( + "Updating mydataset (s3://bucket/path)\n" + "M\tbaz\n" + "A\tfoobar\n" + "D\tfoo\n" + ).expandtabs(8), + }, + ), + ], +) +def test_update(dvc, capsys, mocker, spec, old_lock, new_lock, expected_outputs, lock): + if lock == "missing": + new_lock = spec | old_lock | new_lock + old_lock = None + elif lock == "unchanged": + old_lock = new_lock = spec | old_lock + else: + old_lock = spec | old_lock + new_lock = old_lock | new_lock + + old = dvc.datasets._build_dataset("dvc.yaml", spec, old_lock) + new = dvc.datasets._build_dataset("dvc.yaml", spec, new_lock) + + m = mocker.patch("dvc.repo.datasets.Datasets.update", return_value=(old, new)) + assert main(["dataset", "update", spec["name"]]) == 0 + out, err = capsys.readouterr() + assert out == expected_outputs[lock] + assert not err + + m.assert_called_once() diff --git a/tests/unit/test_dataset.py b/tests/unit/test_dataset.py new file mode 100644 index 000000000..22beaf581 --- /dev/null +++ b/tests/unit/test_dataset.py @@ -0,0 +1,45 @@ +import pytest + +from dvc.repo.datasets import parse_url_and_type + + +@pytest.mark.parametrize( + "url,expected", + [ + ("s3://bucket/path", {"type": "url", "url": "s3://bucket/path"}), + ("gs://bucket/path", {"type": "url", "url": "gs://bucket/path"}), + ("azure://container/path", {"type": "url", "url": "azure://container/path"}), + ( + "remote://remote_name/path", + {"type": "url", "url": "remote://remote_name/path"}, + ), + ("dvcx://dataset_name", {"type": "dvcx", "url": "dvcx://dataset_name"}), + ( + "dvc+file:///home/user/repository", + { + "type": "dvc", + "url": "file:///home/user/repository", + }, + ), + ( + "dvc://git@github.com:iterative/example-get-started.git", + {"type": "dvc", "url": "git@github.com:iterative/example-get-started.git"}, + ), + ( + "dvc+https://github.com/iterative/example-get-started.git", + { + "type": "dvc", + "url": "https://github.com/iterative/example-get-started.git", + }, + ), + ( + "dvc+ssh://github.com/iterative/example-get-started.git", + { + "type": "dvc", + "url": "ssh://github.com/iterative/example-get-started.git", + }, + ), + ], +) +def test_url_parsing(url, expected): + assert parse_url_and_type(url) == expected