From 3767b03784a80fd229ad49f5b57cee32bc26932c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniele=20Trifir=C3=B2?= <daniele@iterative.ai>
Date: Fri, 11 Feb 2022 18:23:26 +0100
Subject: [PATCH] test_data_cloud: also check for the correct number of jobs

---
 tests/func/test_data_cloud.py | 46 +++++++++++++++++++++++++++++++++--
 1 file changed, 44 insertions(+), 2 deletions(-)

diff --git a/tests/func/test_data_cloud.py b/tests/func/test_data_cloud.py
index 9bd56f460..eea0e549b 100644
--- a/tests/func/test_data_cloud.py
+++ b/tests/func/test_data_cloud.py
@@ -7,7 +7,9 @@ from flaky.flaky_decorator import flaky
 
 import dvc as dvc_module
 from dvc.cli import main
+from dvc.data.db.local import LocalObjectDB
 from dvc.external_repo import clean_repos
+from dvc.objects.db import ObjectDB
 from dvc.stage.exceptions import StageNotFound
 from dvc.testing.test_remote import (  # noqa, pylint: disable=unused-import
     TestRemote,
@@ -15,8 +17,9 @@ from dvc.testing.test_remote import (  # noqa, pylint: disable=unused-import
 from dvc.utils.fs import remove
 
 
-def test_cloud_cli(tmp_dir, dvc, remote):
-    args = ["-v", "-j", "2"]
+def test_cloud_cli(tmp_dir, dvc, remote, mocker):
+    jobs = 2
+    args = ["-v", "-j", str(jobs)]
 
     (stage,) = tmp_dir.dvc_gen("foo", "foo")
     cache = stage.outs[0].cache_path
@@ -34,18 +37,32 @@ def test_cloud_cli(tmp_dir, dvc, remote):
     cache_dir = stage_dir.outs[0].cache_path
 
     # FIXME check status output
+    hashes_exist = mocker.spy(LocalObjectDB, "hashes_exist")
 
     assert main(["push"] + args) == 0
     assert os.path.exists(cache)
     assert os.path.isfile(cache)
     assert os.path.isfile(cache_dir)
+    assert hashes_exist.called
+    assert all(
+        _kwargs["jobs"] == jobs
+        for (_args, _kwargs) in hashes_exist.call_args_list
+    )
 
     remove(dvc.odb.local.cache_dir)
+    hashes_exist.reset_mock()
 
     assert main(["fetch"] + args) == 0
     assert os.path.exists(cache)
     assert os.path.isfile(cache)
     assert os.path.isfile(cache_dir)
+    assert hashes_exist.called
+    assert all(
+        _kwargs["jobs"] == jobs
+        for (_args, _kwargs) in hashes_exist.call_args_list
+    )
+
+    hashes_exist.reset_mock()
 
     assert main(["pull"] + args) == 0
     assert os.path.exists(cache)
@@ -53,6 +70,11 @@ def test_cloud_cli(tmp_dir, dvc, remote):
     assert os.path.isfile(cache_dir)
     assert os.path.isfile("foo")
     assert os.path.isdir("data_dir")
+    assert hashes_exist.called
+    assert all(
+        _kwargs["jobs"] == jobs
+        for (_args, _kwargs) in hashes_exist.call_args_list
+    )
 
     with open(cache, encoding="utf-8") as fd:
         assert fd.read() == "foo"
@@ -62,18 +84,38 @@ def test_cloud_cli(tmp_dir, dvc, remote):
     if remote.url.startswith("http"):
         return
 
+    hashes_exist.reset_mock()
+
+    _list_hashes_traverse = mocker.spy(ObjectDB, "_list_hashes_traverse")
     # NOTE: check if remote gc works correctly on directories
     assert main(["gc", "-cw", "-f"] + args) == 0
+    assert _list_hashes_traverse.called
+    assert all(
+        _kwargs["jobs"] == 2
+        for (_args, _kwargs) in hashes_exist.call_args_list
+    )
     shutil.move(dvc.odb.local.cache_dir, dvc.odb.local.cache_dir + ".back")
 
     assert main(["fetch"] + args) == 0
 
+    assert hashes_exist.called
+    assert all(
+        _kwargs["jobs"] == jobs
+        for (_args, _kwargs) in hashes_exist.call_args_list
+    )
+
+    hashes_exist.reset_mock()
     assert main(["pull", "-f"] + args) == 0
     assert os.path.exists(cache)
     assert os.path.isfile(cache)
     assert os.path.isfile(cache_dir)
     assert os.path.isfile("foo")
     assert os.path.isdir("data_dir")
+    assert hashes_exist.called
+    assert all(
+        _kwargs["jobs"] == jobs
+        for (_args, _kwargs) in hashes_exist.call_args_list
+    )
 
 
 def test_data_cloud_error_cli(dvc):