Repository: iterative/dvc
Pull request #10321 (2024-02-26 15:51:59)

tests: add tests for datasets
Changes in the code are cosmetic changes, or done to make it easier to mock in the tests or make assertions easily. 


Patch:
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)


Test patch:
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