From 7c7e73e3d22fe0ca3ababa4a750a4f05790c3c29 Mon Sep 17 00:00:00 2001
From: Saugat Pachhai <suagatchhetri@outlook.com>
Date: Wed, 31 Mar 2021 08:31:27 -0700
Subject: [PATCH] status: collect stages only once (#5741)

This is achieved by using `lock_repo`
context manager that can be used in a recursive depth
that only resets stages when it has left any lock_repo
blocks.

Also added some ui related tests for the status which I found
missing
---
 dvc/command/status.py             | 42 ++++++++++++++++------------
 tests/unit/command/test_status.py | 46 +++++++++++++++++++++++++++++--
 2 files changed, 68 insertions(+), 20 deletions(-)

diff --git a/dvc/command/status.py b/dvc/command/status.py
index 616d22376..741c59710 100644
--- a/dvc/command/status.py
+++ b/dvc/command/status.py
@@ -44,21 +44,27 @@ class CmdDataStatus(CmdDataBase):
                 self._show(value, indent + 1)
 
     def run(self):
+        from dvc.repo import lock_repo
         from dvc.ui import ui
 
         indent = 1 if self.args.cloud else 0
-        try:
-            st = self.repo.status(
-                targets=self.args.targets,
-                jobs=self.args.jobs,
-                cloud=self.args.cloud,
-                remote=self.args.remote,
-                all_branches=self.args.all_branches,
-                all_tags=self.args.all_tags,
-                all_commits=self.args.all_commits,
-                with_deps=self.args.with_deps,
-                recursive=self.args.recursive,
-            )
+
+        with lock_repo(self.repo):
+            try:
+                st = self.repo.status(
+                    targets=self.args.targets,
+                    jobs=self.args.jobs,
+                    cloud=self.args.cloud,
+                    remote=self.args.remote,
+                    all_branches=self.args.all_branches,
+                    all_tags=self.args.all_tags,
+                    all_commits=self.args.all_commits,
+                    with_deps=self.args.with_deps,
+                    recursive=self.args.recursive,
+                )
+            except DvcException:
+                logger.exception("")
+                return 1
 
             if self.args.quiet:
                 return bool(st)
@@ -67,9 +73,14 @@ class CmdDataStatus(CmdDataBase):
                 import json
 
                 logger.info(json.dumps(st))
-            elif st:
+                return 0
+
+            if st:
                 self._show(st, indent)
-            elif not self.repo.stages:
+                return 0
+
+            # additional hints for the user
+            if not self.repo.stages:
                 ui.write(self.EMPTY_PROJECT_MSG)
             elif self.args.cloud or self.args.remote:
                 remote = self.args.remote or self.repo.config["core"].get(
@@ -79,7 +90,4 @@ class CmdDataStatus(CmdDataBase):
             else:
                 ui.write(self.UP_TO_DATE_MSG)
 
-        except DvcException:
-            logger.exception("")
-            return 1
         return 0
diff --git a/tests/unit/command/test_status.py b/tests/unit/command/test_status.py
index 575bf3b21..d6a370270 100644
--- a/tests/unit/command/test_status.py
+++ b/tests/unit/command/test_status.py
@@ -45,7 +45,7 @@ def test_cloud_status(tmp_dir, dvc, mocker):
 
 
 @pytest.mark.parametrize("status", [{}, {"a": "b", "c": [1, 2, 3]}, [1, 2, 3]])
-def test_status_show_json(mocker, caplog, status):
+def test_status_show_json(dvc, mocker, caplog, status):
     cli_args = parse_args(["status", "--show-json"])
     assert cli_args.func == CmdDataStatus
 
@@ -60,7 +60,7 @@ def test_status_show_json(mocker, caplog, status):
 @pytest.mark.parametrize(
     "status, ret", [({}, 0), ({"a": "b", "c": [1, 2, 3]}, 1), ([1, 2, 3], 1)]
 )
-def test_status_quiet(mocker, caplog, capsys, status, ret):
+def test_status_quiet(dvc, mocker, caplog, capsys, status, ret):
     cli_args = parse_args(["status", "-q"])
     assert cli_args.func == CmdDataStatus
 
@@ -71,5 +71,45 @@ def test_status_quiet(mocker, caplog, capsys, status, ret):
     assert cmd.run() == ret
     assert not caplog.messages
     captured = capsys.readouterr()
-    assert not captured.err
     assert not captured.out
+
+
+def test_status_empty(dvc, mocker, capsys):
+    cli_args = parse_args(["status"])
+    assert cli_args.func == CmdDataStatus
+
+    cmd = cli_args.func(cli_args)
+
+    spy = mocker.spy(cmd.repo.stage, "collect_repo")
+
+    assert cmd.run() == 0
+
+    captured = capsys.readouterr()
+    assert "no data or pipelines tracked" in captured.out
+    # stages should only be collected once
+    assert spy.call_count == 1
+
+
+@pytest.mark.parametrize(
+    "cloud_opts, expected_message",
+    [
+        (["--cloud"], "Cache and remote 'default' are in sync"),
+        (["--remote", "remote1"], "Cache and remote 'remote1' are in sync"),
+        ([], "Data and pipelines are up to date"),
+    ],
+)
+def test_status_up_to_date(dvc, mocker, capsys, cloud_opts, expected_message):
+    cli_args = parse_args(["status", *cloud_opts])
+    assert cli_args.func == CmdDataStatus
+
+    cmd = cli_args.func(cli_args)
+
+    mocker.patch.dict(cmd.repo.config, {"core": {"remote": "default"}})
+    mocker.patch.object(cmd.repo, "status", autospec=True, return_value={})
+    mocker.patch.object(
+        cmd.repo.stage, "collect_repo", return_value=[object()], autospec=True
+    )
+
+    assert cmd.run() == 0
+    captured = capsys.readouterr()
+    assert expected_message in captured.out