From a86ba0092793a362760dac52c59fb0ade478b836 Mon Sep 17 00:00:00 2001 From: pared Date: Tue, 5 Mar 2019 13:44:00 +0100 Subject: [PATCH] utils: translate yaml file loading exception to dvc one --- dvc/exceptions.py | 9 +++++++++ dvc/utils/__init__.py | 9 ++++++++- tests/test_add.py | 27 +++++++++++++++++++++++++++ tests/test_data_cloud.py | 9 +++------ tests/test_import.py | 10 +++------- tests/test_logger.py | 9 +++++---- tests/utils/__init__.py | 7 +++++++ tests/utils/logger.py | 15 +++++++++++++++ 8 files changed, 77 insertions(+), 18 deletions(-) diff --git a/dvc/exceptions.py b/dvc/exceptions.py index b31eef8b5..b225d61f8 100644 --- a/dvc/exceptions.py +++ b/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) diff --git a/dvc/utils/__init__.py b/dvc/utils/__init__.py index 13c501c3b..b7b02d4ad 100644 --- a/dvc/utils/__init__.py +++ b/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) diff --git a/tests/test_add.py b/tests/test_add.py index 89fead9d8..af2ee388d 100644 --- a/tests/test_add.py +++ b/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(), + ) diff --git a/tests/test_data_cloud.py b/tests/test_data_cloud.py index 6592e833d..6eb826931 100644 --- a/tests/test_data_cloud.py +++ b/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): diff --git a/tests/test_import.py b/tests/test_import.py index ff3305121..61770c2a2 100644 --- a/tests/test_import.py +++ b/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): diff --git a/tests/test_logger.py b/tests/test_logger.py index 8d1628ae4..2e8a4e36c 100644 --- a/tests/test_logger.py +++ b/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): diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py index 8a9f854a7..f452a1c31 100644 --- a/tests/utils/__init__.py +++ b/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() diff --git a/tests/utils/logger.py b/tests/utils/logger.py index cc34b1f46..ed127e778 100644 --- a/tests/utils/logger.py +++ b/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()