tests: add tests for datasets (#10321)

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()

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://[email protected]:iterative/example-get-started.git",
+            {"type": "dvc", "url": "[email protected]: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