tests/repo: move .collect to unit

tests/func/test_repo.py

@@ -1,52 +1,10 @@
 from dvc.ignore import DvcIgnore
 from dvc.main import main
 from dvc.repo import Repo
-from dvc.scm.git import GitTree
-from dvc.scm.tree import WorkingTree
 from dvc.stage import Stage
 from tests.basic_env import TestDvcGit
 
 
-class TestCollect(TestDvcGit):
-    def setUp(self):
-        super(TestCollect, self).setUp()
-        self.dvc.add(self.FOO)
-        self.dvc.run(
-            deps=[self.FOO],
-            outs=[self.BAR],
-            cmd="python code.py {} {}".format(self.FOO, self.BAR),
-        )
-        self.dvc.scm.add([".gitignore", self.FOO + ".dvc", self.BAR + ".dvc"])
-        self.dvc.scm.commit("foo.dvc and bar.dvc")
-        self.dvc.scm.checkout("new_branch", True)
-        self.dvc.run(
-            deps=[self.BAR],
-            outs=["buzz"],
-            cmd="python code.py {} {}".format(self.BAR, "buzz"),
-        )
-        self.dvc.scm.add([".gitignore", "buzz.dvc"])
-        self.dvc.scm.commit("add buzz")
-        self.dvc.scm.checkout("master")
-
-    def _check(self, branch, target, with_deps, expected):
-        if branch:
-            self.dvc.tree = GitTree(self.dvc.scm.repo, branch)
-        else:
-            self.dvc.tree = WorkingTree()
-        result = self.dvc.collect(target + ".dvc", with_deps=with_deps)
-        self.assertEqual([[str(j) for j in i.outs] for i in result], expected)
-        return result
-
-    def test(self):
-        self._check("", self.BAR, True, [[self.FOO], [self.BAR]])
-        self._check("master", self.BAR, True, [[self.FOO], [self.BAR]])
-        self._check(
-            "new_branch", "buzz", True, [[self.FOO], [self.BAR], ["buzz"]]
-        )
-        result = self._check("new_branch", "buzz", False, [["buzz"]])
-        self.assertEqual([str(i) for i in result[0].deps], ["bar"])
-
-
 class TestIgnore(TestDvcGit):
     def _stage_name(self, file):
         return file + Stage.STAGE_FILE_SUFFIX

tests/unit/test_repo.py

@@ -1,3 +1,4 @@
+from dvc.scm.git import GitTree
 from dvc.system import System
 from dvc.utils.compat import fspath
 
@@ -32,3 +33,30 @@ def test_destroy(tmp_dir, dvc):
     assert not System.is_symlink(fspath(tmp_dir / "dir"))
     assert not System.is_symlink(fspath(tmp_dir / "dir" / "file"))
     assert not System.is_symlink(fspath(tmp_dir / "dir" / "subdir" / "file"))
+
+
+def test_collect(tmp_dir, scm, dvc, run_copy):
+    tmp_dir.dvc_gen("foo", "foo")
+    run_copy("foo", "bar")
+    scm.add([".gitignore", "foo.dvc", "bar.dvc"])
+    scm.commit("Add foo and bar")
+
+    scm.checkout("new-branch", create_new=True)
+
+    run_copy("bar", "buzz")
+    scm.add([".gitignore", "buzz.dvc"])
+    scm.commit("Add buzz")
+
+    def collect_outs(*args, **kwargs):
+        return [
+            str(out.path_info)
+            for stage in dvc.collect(*args, **kwargs)
+            for out in stage.outs
+        ]
+
+    assert ["foo", "bar"] == collect_outs("bar.dvc", with_deps=True)
+
+    dvc.tree = GitTree(scm.repo, "new-branch")
+
+    assert ["foo", "bar", "buzz"] == collect_outs("buzz.dvc", with_deps=True)
+    assert ["buzz"] == collect_outs("buzz.dvc", with_deps=False)