From 3741f730e0e8a7b839d8e7690a1cc2b3a342e5f8 Mon Sep 17 00:00:00 2001 From: Saugat Pachhai Date: Mon, 25 Jan 2021 02:28:38 +0545 Subject: [PATCH] show `git-add` hint on `stage:add`, refactor scm_context (#5290) * Refactor scm_context, show git-add hint on stage:add * move tests --- dvc/command/stage.py | 8 +-- dvc/repo/scm_context.py | 27 +++++----- dvc/scm/base.py | 20 +++++++ tests/unit/scm/test_git.py | 6 +++ tests/unit/scm/test_scm.py | 103 +++++++++++++++++++------------------ 5 files changed, 96 insertions(+), 68 deletions(-) diff --git a/dvc/command/stage.py b/dvc/command/stage.py index 7ca1355d7..c52adbae8 100644 --- a/dvc/command/stage.py +++ b/dvc/command/stage.py @@ -107,11 +107,13 @@ class CmdStageList(CmdBase): class CmdStageAdd(CmdBase): def run(self): + repo = self.repo kwargs = vars(self.args) - stage = self.repo.stage.create_from_cli(validate=True, **kwargs) + stage = repo.stage.create_from_cli(validate=True, **kwargs) - stage.ignore_outs() - stage.dump() + with repo.scm.track_file_changes(config=repo.config): + stage.ignore_outs() + stage.dump() return 0 diff --git a/dvc/repo/scm_context.py b/dvc/repo/scm_context.py index ec9585bb7..b0ecc2df5 100644 --- a/dvc/repo/scm_context.py +++ b/dvc/repo/scm_context.py @@ -1,17 +1,16 @@ +from functools import wraps +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from dvc.repo import Repo + + def scm_context(method): - def run(repo, *args, **kw): - try: - result = method(repo, *args, **kw) - repo.scm.reset_ignores() - autostage = repo.config.get("core", {}).get("autostage", False) - if autostage: - repo.scm.track_changed_files() - else: - repo.scm.remind_to_track() - repo.scm.reset_tracked_files() - return result - except Exception: - repo.scm.cleanup_ignores() - raise + @wraps(method) + def run(repo: "Repo", *args, **kw): + scm = repo.scm + + with scm.track_file_changes(config=repo.config): + return method(repo, *args, **kw) return run diff --git a/dvc/scm/base.py b/dvc/scm/base.py index c795f9460..b47b89c59 100644 --- a/dvc/scm/base.py +++ b/dvc/scm/base.py @@ -1,6 +1,9 @@ """Manages source control systems (e.g. Git) in DVC.""" import os +from contextlib import contextmanager + +from funcy import get_in from dvc.exceptions import DvcException @@ -173,6 +176,23 @@ class Base: to track new files """ + @contextmanager + def track_file_changes(self, config=None): + autostage = get_in(config or {}, ["core", "autostage"]) + + try: + yield + except Exception: + self.cleanup_ignores() + raise + + self.reset_ignores() + if autostage: + self.track_changed_files() + else: + self.remind_to_track() + self.reset_tracked_files() + def belongs_to_scm(self, path): """Return boolean whether file belongs to scm""" diff --git a/tests/unit/scm/test_git.py b/tests/unit/scm/test_git.py index 3a5b669aa..b81849e3b 100644 --- a/tests/unit/scm/test_git.py +++ b/tests/unit/scm/test_git.py @@ -498,3 +498,9 @@ def test_reset(tmp_dir, scm, git): staged, unstaged, _ = scm.status() assert len(staged) == 1 assert len(unstaged) == 1 + + +def test_remind_to_track(scm, caplog): + scm.files_to_track = ["fname with spaces.txt", "тест", "foo"] + scm.remind_to_track() + assert "git add 'fname with spaces.txt' 'тест' foo" in caplog.text diff --git a/tests/unit/scm/test_scm.py b/tests/unit/scm/test_scm.py index 1121d4df5..793bf8a1a 100644 --- a/tests/unit/scm/test_scm.py +++ b/tests/unit/scm/test_scm.py @@ -1,60 +1,61 @@ -from unittest import TestCase +import pytest -import mock - -from dvc.repo import Repo from dvc.repo.scm_context import scm_context from dvc.scm import NoSCM -class TestScmContext(TestCase): - def setUp(self): - self.repo_mock = mock.Mock(spec=Repo) - self.scm_mock = mock.Mock(spec=NoSCM) - self.repo_mock.scm = self.scm_mock +def test_scm_context(dvc, mocker): + scm = dvc.scm + m = mocker.patch.object( + scm, "track_file_changes", wraps=scm.track_file_changes + ) - def test_should_successfully_perform_method(self): - method = mock.Mock() - wrapped = scm_context(method) + ret_val = 0 + wrapped = scm_context(mocker.Mock(return_value=ret_val)) - self.repo_mock.configure_mock(config={}) - wrapped(self.repo_mock) - - self.assertEqual(1, method.call_count) - self.assertEqual(1, self.scm_mock.reset_ignores.call_count) - self.assertEqual(1, self.scm_mock.remind_to_track.call_count) - - self.assertEqual(0, self.scm_mock.cleanup_ignores.call_count) - - def test_should_check_autostage(self): - method = mock.Mock() - wrapped = scm_context(method) - - config_autostage_attrs = {"config": {"core": {"autostage": True}}} - self.repo_mock.configure_mock(**config_autostage_attrs) - wrapped(self.repo_mock) - - self.assertEqual(1, method.call_count) - self.assertEqual(1, self.scm_mock.reset_ignores.call_count) - self.assertEqual(1, self.scm_mock.track_changed_files.call_count) - - self.assertEqual(0, self.scm_mock.cleanup_ignores.call_count) - - def test_should_throw_and_cleanup(self): - method = mock.Mock(side_effect=Exception("some problem")) - wrapped = scm_context(method) - - with self.assertRaises(Exception): - wrapped(self.repo_mock) - - self.assertEqual(1, method.call_count) - self.assertEqual(1, self.scm_mock.cleanup_ignores.call_count) - - self.assertEqual(0, self.scm_mock.reset_ignores.call_count) - self.assertEqual(0, self.scm_mock.remind_to_track.call_count) + assert wrapped(dvc) == ret_val + m.assert_called_once_with(config=dvc.config) -def test_remind_to_track(scm, caplog): - scm.files_to_track = ["fname with spaces.txt", "тест", "foo"] - scm.remind_to_track() - assert "git add 'fname with spaces.txt' 'тест' foo" in caplog.text +def test_track_file_changes(mocker): + scm = mocker.Mock(NoSCM) + + with NoSCM.track_file_changes(scm): + pass + + assert scm.reset_ignores.call_count == 1 + assert scm.remind_to_track.call_count == 1 + assert scm.track_changed_files.call_count == 0 + assert scm.cleanup_ignores.call_count == 0 + assert scm.reset_tracked_files.call_count == 1 + + +def test_track_file_changes_autostage(mocker): + scm = mocker.Mock(NoSCM) + + config = {"core": {"autostage": True}} + with NoSCM.track_file_changes(scm, config=config): + pass + + assert scm.track_changed_files.call_count == 1 + assert scm.reset_ignores.call_count == 1 + assert scm.remind_to_track.call_count == 0 + assert scm.cleanup_ignores.call_count == 0 + assert scm.reset_tracked_files.call_count == 1 + + +def test_track_file_changes_throw_and_cleanup(mocker): + scm = mocker.Mock(NoSCM) + + class CustomException(Exception): + pass + + with pytest.raises(CustomException, match="oops"): + with NoSCM.track_file_changes(scm): + raise CustomException("oops") + + assert scm.cleanup_ignores.call_count == 1 + assert scm.reset_ignores.call_count == 0 + assert scm.track_changed_files.call_count == 0 + assert scm.remind_to_track.call_count == 0 + assert scm.reset_tracked_files.call_count == 0