utils: translate yaml file loading exception to dvc one

dvc/exceptions.py

@@ -208,3 +208,12 @@ class NoMetricsError(DvcException):
             "no metric files in this repository. "
             "Use 'dvc metrics add' to add a metric file to track."
         )
+
+
+class StageFileCorruptedError(DvcException):
+    def __init__(self, path):
+        msg = (
+            "Unable to read stage file: {} "
+            "\n Yaml file structure is corrupted".format(path)
+        )
+        super(StageFileCorruptedError, self).__init__(msg)

dvc/utils/__init__.py

@@ -15,6 +15,8 @@ import hashlib
 import nanotime
 import time
 
+from yaml.scanner import ScannerError
+
 LOCAL_CHUNK_SIZE = 1024 * 1024
 LARGE_FILE_SIZE = 1024 * 1024 * 1024
 LARGE_DIR_SIZE = 100
@@ -230,5 +232,10 @@ def current_timestamp():
 
 
 def load_stage_file(path):
+    from dvc.exceptions import StageFileCorruptedError
+
     with open(path, "r") as fobj:
-        return yaml.safe_load(fobj) or {}
+        try:
+            return yaml.safe_load(fobj) or {}
+        except ScannerError:
+            raise StageFileCorruptedError(path)

tests/test_add.py

@@ -1,6 +1,7 @@
 from __future__ import unicode_literals
 
 import os
+from dvc.utils.compat import StringIO
 
 import dvc
 import time
@@ -8,6 +9,7 @@ import shutil
 import filecmp
 import posixpath
 
+from dvc.logger import logger
 from dvc.system import System
 from mock import patch
 
@@ -20,6 +22,7 @@ from dvc.repo import Repo as DvcRepo
 
 from tests.basic_env import TestDvc
 from tests.utils import spy
+from tests.utils.logger import MockLoggerHandlers, ConsoleFontColorsRemover
 
 
 class TestAdd(TestDvc):
@@ -402,3 +405,27 @@ class TestShouldPlaceStageInDataDirIfRepositoryBelowSymlink(TestDvc):
 
             stage_file_path = self.DATA + Stage.STAGE_FILE_SUFFIX
             self.assertTrue(os.path.exists(stage_file_path))
+
+
+class TestShouldThrowProperExceptionOnCorruptedStageFile(TestDvc):
+    def test(self):
+        with MockLoggerHandlers(logger), ConsoleFontColorsRemover():
+            logger.handlers[1].stream = StringIO()
+
+            ret = main(["add", self.FOO])
+            self.assertEqual(0, ret)
+
+            foo_stage = os.path.abspath(self.FOO + Stage.STAGE_FILE_SUFFIX)
+
+            # corrupt stage file
+            with open(foo_stage, "a+") as file:
+                file.write("this will break yaml file structure")
+
+            ret = main(["add", self.BAR])
+            self.assertEqual(1, ret)
+
+            self.assertIn(
+                "Unable to read stage file: {} \n "
+                "Yaml file structure is corrupted".format(foo_stage),
+                logger.handlers[1].stream.getvalue(),
+            )

tests/test_data_cloud.py

@@ -11,7 +11,7 @@ import copy
 
 from dvc.state import State
 from mock import patch
-from tests.utils.logger import MockLoggerHandlers
+from tests.utils.logger import MockLoggerHandlers, ConsoleFontColorsRemover
 
 import dvc.logger as logger
 from dvc.utils.compat import str
@@ -617,8 +617,6 @@ class TestDataCloudErrorCLI(TestDvc):
 
 
 class TestWarnOnOutdatedStage(TestDvc):
-    color_patch = patch.object(logger, "colorize", new=lambda x, color="": x)
-
     def main(self, args):
         ret = main(args)
         self.assertEqual(ret, 0)
@@ -649,9 +647,8 @@ class TestWarnOnOutdatedStage(TestDvc):
             )
 
     def test(self):
-        self.color_patch.start()
-        self._test()
-        self.color_patch.stop()
+        with ConsoleFontColorsRemover():
+            self._test()
 
 
 class TestRecursiveSyncOperations(TestDataCloudBase):

tests/test_import.py

@@ -10,8 +10,7 @@ from dvc.main import main
 from mock import patch, mock_open, call
 from tests.basic_env import TestDvc
 from tests.utils.httpd import StaticFileServer
-from tests.utils.logger import MockLoggerHandlers
-
+from tests.utils.logger import MockLoggerHandlers, ConsoleFontColorsRemover
 
 from dvc.utils.compat import StringIO
 
@@ -46,12 +45,9 @@ class TestDefaultOutput(TestDvc):
 
 
 class TestFailedImportMessage(TestDvc):
-    color_patch = patch.object(logger, "colorize", new=lambda x, color="": x)
-
     def test(self):
-        self.color_patch.start()
-        self._test()
-        self.color_patch.stop()
+        with ConsoleFontColorsRemover():
+            self._test()
 
     @patch("dvc.command.imp.urlparse")
     def _test(self, imp_urlparse_patch):

tests/test_logger.py

@@ -1,5 +1,4 @@
 import traceback
-from mock import patch
 from unittest import TestCase
 
 import logging
@@ -7,11 +6,11 @@ import dvc.logger as logger
 from dvc.command.base import CmdBase
 from dvc.exceptions import DvcException
 from dvc.utils.compat import StringIO
+from tests.utils.logger import ConsoleFontColorsRemover
 
 
 class TestLogger(TestCase):
     handlers = logger.logger.handlers
-    color_patch = patch.object(logger, "colorize", new=lambda x, color="": x)
 
     def setUp(self):
         logger.logger.handlers = [
@@ -21,12 +20,14 @@ class TestLogger(TestCase):
         logger.logger.handlers[0].stream = StringIO()
         logger.logger.handlers[1].stream = StringIO()
         logger.set_default_level()
-        self.color_patch.start()
+
+        self.consoleColorRemover = ConsoleFontColorsRemover()
+        self.consoleColorRemover.__enter__()
 
     def tearDown(self):
         logger.logger.handlers = self.handlers
         logger.be_verbose()
-        self.color_patch.stop()
+        self.consoleColorRemover.__exit__()
 
     @property
     def stdout(self):

tests/utils/__init__.py

@@ -1,4 +1,6 @@
 import os
+
+from dvc.scm import Git
 from mock import MagicMock
 from contextlib import contextmanager
 
@@ -14,6 +16,11 @@ def spy(method_to_decorate):
     return wrapper
 
 
+def get_gitignore_content():
+    with open(Git.GITIGNORE, "r") as gitignore:
+        return gitignore.readlines()
+
+
 @contextmanager
 def cd(newdir):
     prevdir = os.getcwd()

tests/utils/logger.py

@@ -1,4 +1,6 @@
 import logging
+from mock import patch
+import dvc.logger as logger
 
 
 class MockLoggerHandlers(object):
@@ -18,3 +20,16 @@ class MockLoggerHandlers(object):
         for h in self._logger.handlers:
             h.close()
         self._logger.handlers = self._handlers
+
+
+class ConsoleFontColorsRemover(object):
+    def __init__(self):
+        self.color_patch = patch.object(
+            logger, "colorize", new=lambda x, color="": x
+        )
+
+    def __enter__(self):
+        self.color_patch.start()
+
+    def __exit__(self, exc_type=None, exc_val=None, exc_tb=None):
+        self.color_patch.stop()