Add RunContext helper for cli tests

Single context manager that includes exit code and output streams.

Use new RunContext throughout test_cli.

Largely non-functional change, saving some repetition of setup.

Also improve some failures by bundling multiple assertions into one.

tests/test_cli.py

@@ -32,6 +32,29 @@ from yamllint import cli
 from yamllint import config
 
 
+class RunContext(object):
+    """Context manager for ``cli.run()`` to capture exit code and streams."""
+
+    def __init__(self, case):
+        self.stdout = self.stderr = None
+        self._raises_ctx = case.assertRaises(SystemExit)
+
+    def __enter__(self):
+        self._raises_ctx.__enter__()
+        sys.stdout = self.outstream = StringIO()
+        sys.stderr = self.errstream = StringIO()
+        return self
+
+    def __exit__(self, *exc_info):
+        self.stdout, sys.stdout = self.outstream.getvalue(), sys.__stdout__
+        self.stderr, sys.stderr = self.errstream.getvalue(), sys.__stderr__
+        return self._raises_ctx.__exit__(*exc_info)
+
+    @property
+    def returncode(self):
+        return self._raises_ctx.exception.code
+
+
 class CommandLineTestCase(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
@@ -174,201 +197,144 @@ class CommandLineTestCase(unittest.TestCase):
         )
 
     def test_run_with_bad_arguments(self):
-        sys.stdout, sys.stderr = StringIO(), StringIO()
-        with self.assertRaises(SystemExit) as ctx:
+        with RunContext(self) as ctx:
             cli.run(())
+        self.assertNotEqual(ctx.returncode, 0)
+        self.assertEqual(ctx.stdout, '')
+        self.assertRegexpMatches(ctx.stderr, r'^usage')
 
-        self.assertNotEqual(ctx.exception.code, 0)
-
-        out, err = sys.stdout.getvalue(), sys.stderr.getvalue()
-        self.assertEqual(out, '')
-        self.assertRegexpMatches(err, r'^usage')
-
-        sys.stdout, sys.stderr = StringIO(), StringIO()
-        with self.assertRaises(SystemExit) as ctx:
+        with RunContext(self) as ctx:
             cli.run(('--unknown-arg', ))
+        self.assertNotEqual(ctx.returncode, 0)
+        self.assertEqual(ctx.stdout, '')
+        self.assertRegexpMatches(ctx.stderr, r'^usage')
 
-        self.assertNotEqual(ctx.exception.code, 0)
-
-        out, err = sys.stdout.getvalue(), sys.stderr.getvalue()
-        self.assertEqual(out, '')
-        self.assertRegexpMatches(err, r'^usage')
-
-        sys.stdout, sys.stderr = StringIO(), StringIO()
-        with self.assertRaises(SystemExit) as ctx:
+        with RunContext(self) as ctx:
             cli.run(('-c', './conf.yaml', '-d', 'relaxed', 'file'))
-
-        self.assertNotEqual(ctx.exception.code, 0)
-
-        out, err = sys.stdout.getvalue(), sys.stderr.getvalue()
-        self.assertEqual(out, '')
+        self.assertNotEqual(ctx.returncode, 0)
+        self.assertEqual(ctx.stdout, '')
         self.assertRegexpMatches(
-            err.splitlines()[-1],
+            ctx.stderr.splitlines()[-1],
             r'^yamllint: error: argument -d\/--config-data: '
             r'not allowed with argument -c\/--config-file$'
         )
 
         # checks if reading from stdin and files are mutually exclusive
-        sys.stdout, sys.stderr = StringIO(), StringIO()
-        with self.assertRaises(SystemExit) as ctx:
+        with RunContext(self) as ctx:
             cli.run(('-', 'file'))
-
-        self.assertNotEqual(ctx.exception.code, 0)
-
-        out, err = sys.stdout.getvalue(), sys.stderr.getvalue()
-        self.assertEqual(out, '')
-        self.assertRegexpMatches(err, r'^usage')
+        self.assertNotEqual(ctx.returncode, 0)
+        self.assertEqual(ctx.stdout, '')
+        self.assertRegexpMatches(ctx.stderr, r'^usage')
 
     def test_run_with_bad_config(self):
-        sys.stdout, sys.stderr = StringIO(), StringIO()
-        with self.assertRaises(SystemExit) as ctx:
+        with RunContext(self) as ctx:
             cli.run(('-d', 'rules: {a: b}', 'file'))
-
-        self.assertEqual(ctx.exception.code, -1)
-
-        out, err = sys.stdout.getvalue(), sys.stderr.getvalue()
-        self.assertEqual(out, '')
-        self.assertRegexpMatches(err, r'^invalid config: no such rule')
+        self.assertEqual(ctx.returncode, -1)
+        self.assertEqual(ctx.stdout, '')
+        self.assertRegexpMatches(ctx.stderr, r'^invalid config: no such rule')
 
     def test_run_with_empty_config(self):
-        sys.stdout, sys.stderr = StringIO(), StringIO()
-        with self.assertRaises(SystemExit) as ctx:
+        with RunContext(self) as ctx:
             cli.run(('-d', '', 'file'))
-
-        self.assertEqual(ctx.exception.code, -1)
-
-        out, err = sys.stdout.getvalue(), sys.stderr.getvalue()
-        self.assertEqual(out, '')
-        self.assertRegexpMatches(err, r'^invalid config: not a dict')
+        self.assertEqual(ctx.returncode, -1)
+        self.assertEqual(ctx.stdout, '')
+        self.assertRegexpMatches(ctx.stderr, r'^invalid config: not a dict')
 
     def test_run_with_config_file(self):
         with open(os.path.join(self.wd, 'config'), 'w') as f:
             f.write('rules: {trailing-spaces: disable}')
 
-        with self.assertRaises(SystemExit) as ctx:
+        with RunContext(self) as ctx:
             cli.run(('-c', f.name, os.path.join(self.wd, 'a.yaml')))
-        self.assertEqual(ctx.exception.code, 0)
+        self.assertEqual(ctx.returncode, 0)
 
         with open(os.path.join(self.wd, 'config'), 'w') as f:
             f.write('rules: {trailing-spaces: enable}')
 
-        with self.assertRaises(SystemExit) as ctx:
+        with RunContext(self) as ctx:
             cli.run(('-c', f.name, os.path.join(self.wd, 'a.yaml')))
-        self.assertEqual(ctx.exception.code, 1)
+        self.assertEqual(ctx.returncode, 1)
 
     def test_run_with_user_global_config_file(self):
         home = os.path.join(self.wd, 'fake-home')
-        os.mkdir(home)
-        dir = os.path.join(home, '.config')
-        os.mkdir(dir)
-        dir = os.path.join(dir, 'yamllint')
-        os.mkdir(dir)
+        dir = os.path.join(home, '.config', 'yamllint')
+        os.makedirs(dir)
         config = os.path.join(dir, 'config')
 
-        temp = os.environ['HOME']
+        self.addCleanup(os.environ.update, HOME=os.environ['HOME'])
         os.environ['HOME'] = home
 
         with open(config, 'w') as f:
             f.write('rules: {trailing-spaces: disable}')
 
-        with self.assertRaises(SystemExit) as ctx:
+        with RunContext(self) as ctx:
             cli.run((os.path.join(self.wd, 'a.yaml'), ))
-        self.assertEqual(ctx.exception.code, 0)
+        self.assertEqual(ctx.returncode, 0)
 
         with open(config, 'w') as f:
             f.write('rules: {trailing-spaces: enable}')
 
-        with self.assertRaises(SystemExit) as ctx:
+        with RunContext(self) as ctx:
             cli.run((os.path.join(self.wd, 'a.yaml'), ))
-        self.assertEqual(ctx.exception.code, 1)
-
-        os.environ['HOME'] = temp
+        self.assertEqual(ctx.returncode, 1)
 
     def test_run_version(self):
-        sys.stdout, sys.stderr = StringIO(), StringIO()
-        with self.assertRaises(SystemExit) as ctx:
+        with RunContext(self) as ctx:
             cli.run(('--version', ))
-
-        self.assertEqual(ctx.exception.code, 0)
-
-        out, err = sys.stdout.getvalue(), sys.stderr.getvalue()
-        self.assertRegexpMatches(out + err, r'yamllint \d+\.\d+')
+        self.assertEqual(ctx.returncode, 0)
+        self.assertRegexpMatches(ctx.stdout + ctx.stderr, r'yamllint \d+\.\d+')
 
     def test_run_non_existing_file(self):
-        file = os.path.join(self.wd, 'i-do-not-exist.yaml')
-
-        sys.stdout, sys.stderr = StringIO(), StringIO()
-        with self.assertRaises(SystemExit) as ctx:
-            cli.run(('-f', 'parsable', file))
-
-        self.assertEqual(ctx.exception.code, -1)
+        path = os.path.join(self.wd, 'i-do-not-exist.yaml')
 
-        out, err = sys.stdout.getvalue(), sys.stderr.getvalue()
-        self.assertEqual(out, '')
-        self.assertRegexpMatches(err, r'No such file or directory')
+        with RunContext(self) as ctx:
+            cli.run(('-f', 'parsable', path))
+        self.assertEqual(ctx.returncode, -1)
+        self.assertEqual(ctx.stdout, '')
+        self.assertRegexpMatches(ctx.stderr, r'No such file or directory')
 
     def test_run_one_problem_file(self):
-        file = os.path.join(self.wd, 'a.yaml')
+        path = os.path.join(self.wd, 'a.yaml')
 
-        sys.stdout, sys.stderr = StringIO(), StringIO()
-        with self.assertRaises(SystemExit) as ctx:
-            cli.run(('-f', 'parsable', file))
-
-        self.assertEqual(ctx.exception.code, 1)
-
-        out, err = sys.stdout.getvalue(), sys.stderr.getvalue()
-        self.assertEqual(out, (
+        with RunContext(self) as ctx:
+            cli.run(('-f', 'parsable', path))
+        self.assertEqual(ctx.returncode, 1)
+        self.assertEqual(ctx.stdout, (
             '%s:2:4: [error] trailing spaces (trailing-spaces)\n'
             '%s:3:4: [error] no new line character at the end of file '
-            '(new-line-at-end-of-file)\n') % (file, file))
-        self.assertEqual(err, '')
+            '(new-line-at-end-of-file)\n' % (path, path)))
+        self.assertEqual(ctx.stderr, '')
 
     def test_run_one_warning(self):
-        file = os.path.join(self.wd, 'warn.yaml')
-
-        sys.stdout, sys.stderr = StringIO(), StringIO()
-        with self.assertRaises(SystemExit) as ctx:
-            cli.run(('-f', 'parsable', file))
+        path = os.path.join(self.wd, 'warn.yaml')
 
-        self.assertEqual(ctx.exception.code, 0)
+        with RunContext(self) as ctx:
+            cli.run(('-f', 'parsable', path))
+        self.assertEqual(ctx.returncode, 0)
 
     def test_run_warning_in_strict_mode(self):
-        file = os.path.join(self.wd, 'warn.yaml')
+        path = os.path.join(self.wd, 'warn.yaml')
 
-        sys.stdout, sys.stderr = StringIO(), StringIO()
-        with self.assertRaises(SystemExit) as ctx:
-            cli.run(('-f', 'parsable', '--strict', file))
-
-        self.assertEqual(ctx.exception.code, 2)
+        with RunContext(self) as ctx:
+            cli.run(('-f', 'parsable', '--strict', path))
+        self.assertEqual(ctx.returncode, 2)
 
     def test_run_one_ok_file(self):
-        file = os.path.join(self.wd, 'sub', 'ok.yaml')
-
-        sys.stdout, sys.stderr = StringIO(), StringIO()
-        with self.assertRaises(SystemExit) as ctx:
-            cli.run(('-f', 'parsable', file))
-
-        self.assertEqual(ctx.exception.code, 0)
+        path = os.path.join(self.wd, 'sub', 'ok.yaml')
 
-        out, err = sys.stdout.getvalue(), sys.stderr.getvalue()
-        self.assertEqual(out, '')
-        self.assertEqual(err, '')
+        with RunContext(self) as ctx:
+            cli.run(('-f', 'parsable', path))
+        self.assertEqual((ctx.returncode, ctx.stdout, ctx.stderr), (0, '', ''))
 
     def test_run_empty_file(self):
-        file = os.path.join(self.wd, 'empty.yml')
+        path = os.path.join(self.wd, 'empty.yml')
 
-        sys.stdout, sys.stderr = StringIO(), StringIO()
-        with self.assertRaises(SystemExit) as ctx:
-            cli.run(('-f', 'parsable', file))
-
-        self.assertEqual(ctx.exception.code, 0)
-
-        out, err = sys.stdout.getvalue(), sys.stderr.getvalue()
-        self.assertEqual(out, '')
-        self.assertEqual(err, '')
+        with RunContext(self) as ctx:
+            cli.run(('-f', 'parsable', path))
+        self.assertEqual((ctx.returncode, ctx.stdout, ctx.stderr), (0, '', ''))
 
     def test_run_non_ascii_file(self):
-        file = os.path.join(self.wd, 'non-ascii', 'éçäγλνπ¥', 'utf-8')
+        path = os.path.join(self.wd, 'non-ascii', 'éçäγλνπ¥', 'utf-8')
 
         # Make sure the default localization conditions on this "system"
         # support UTF-8 encoding.
@@ -377,63 +343,46 @@ class CommandLineTestCase(unittest.TestCase):
             locale.setlocale(locale.LC_ALL, 'C.UTF-8')
         except locale.Error:
             locale.setlocale(locale.LC_ALL, 'en_US.UTF-8')
+        self.addCleanup(locale.setlocale, locale.LC_ALL, loc)
 
-        sys.stdout, sys.stderr = StringIO(), StringIO()
-        with self.assertRaises(SystemExit) as ctx:
-            cli.run(('-f', 'parsable', file))
-
-        locale.setlocale(locale.LC_ALL, loc)
-
-        self.assertEqual(ctx.exception.code, 0)
-
-        out, err = sys.stdout.getvalue(), sys.stderr.getvalue()
-        self.assertEqual(out, '')
-        self.assertEqual(err, '')
+        with RunContext(self) as ctx:
+            cli.run(('-f', 'parsable', path))
+        self.assertEqual((ctx.returncode, ctx.stdout, ctx.stderr), (0, '', ''))
 
     def test_run_multiple_files(self):
         items = [os.path.join(self.wd, 'empty.yml'),
                  os.path.join(self.wd, 's')]
-        file = items[1] + '/s/s/s/s/s/s/s/s/s/s/s/s/s/s/file.yaml'
+        path = items[1] + '/s/s/s/s/s/s/s/s/s/s/s/s/s/s/file.yaml'
 
-        sys.stdout, sys.stderr = StringIO(), StringIO()
-        with self.assertRaises(SystemExit) as ctx:
+        with RunContext(self) as ctx:
             cli.run(['-f', 'parsable'] + items)
-
-        self.assertEqual(ctx.exception.code, 1)
-
-        out, err = sys.stdout.getvalue(), sys.stderr.getvalue()
-        self.assertEqual(out, (
+        self.assertEqual((ctx.returncode, ctx.stderr), (1, ''))
+        self.assertEqual(ctx.stdout, (
             '%s:3:1: [error] duplication of key "key" in mapping '
-            '(key-duplicates)\n') % file)
-        self.assertEqual(err, '')
+            '(key-duplicates)\n') % path)
 
     def test_run_piped_output_nocolor(self):
-        file = os.path.join(self.wd, 'a.yaml')
-
-        sys.stdout, sys.stderr = StringIO(), StringIO()
-        with self.assertRaises(SystemExit) as ctx:
-            cli.run((file, ))
+        path = os.path.join(self.wd, 'a.yaml')
 
-        self.assertEqual(ctx.exception.code, 1)
-
-        out, err = sys.stdout.getvalue(), sys.stderr.getvalue()
-        self.assertEqual(out, (
+        with RunContext(self) as ctx:
+            cli.run((path, ))
+        self.assertEqual((ctx.returncode, ctx.stderr), (1, ''))
+        self.assertEqual(ctx.stdout, (
             '%s\n'
             '  2:4       error    trailing spaces  (trailing-spaces)\n'
             '  3:4       error    no new line character at the end of file  '
             '(new-line-at-end-of-file)\n'
-            '\n' % file))
-        self.assertEqual(err, '')
+            '\n' % path))
 
     def test_run_default_format_output_in_tty(self):
-        file = os.path.join(self.wd, 'a.yaml')
+        path = os.path.join(self.wd, 'a.yaml')
 
         # Create a pseudo-TTY and redirect stdout to it
         master, slave = pty.openpty()
         sys.stdout = sys.stderr = os.fdopen(slave, 'w')
 
         with self.assertRaises(SystemExit) as ctx:
-            cli.run((file, ))
+            cli.run((path, ))
         sys.stdout.flush()
 
         self.assertEqual(ctx.exception.code, 1)
@@ -456,114 +405,91 @@ class CommandLineTestCase(unittest.TestCase):
             '  \033[2m3:4\033[0m       \033[31merror\033[0m    '
             'no new line character at the end of file  '
             '\033[2m(new-line-at-end-of-file)\033[0m\n'
-            '\n' % file))
+            '\n' % path))
 
     def test_run_default_format_output_without_tty(self):
-        file = os.path.join(self.wd, 'a.yaml')
-
-        sys.stdout, sys.stderr = StringIO(), StringIO()
-        with self.assertRaises(SystemExit) as ctx:
-            cli.run((file, ))
-
-        self.assertEqual(ctx.exception.code, 1)
+        path = os.path.join(self.wd, 'a.yaml')
 
-        out, err = sys.stdout.getvalue(), sys.stderr.getvalue()
-        self.assertEqual(out, (
+        with RunContext(self) as ctx:
+            cli.run((path, ))
+        expected_out = (
             '%s\n'
             '  2:4       error    trailing spaces  (trailing-spaces)\n'
             '  3:4       error    no new line character at the end of file  '
             '(new-line-at-end-of-file)\n'
-            '\n' % file))
-        self.assertEqual(err, '')
+            '\n' % path)
+        self.assertEqual(
+            (ctx.returncode, ctx.stdout, ctx.stderr), (1, expected_out, ''))
 
     def test_run_auto_output_without_tty_output(self):
-        file = os.path.join(self.wd, 'a.yaml')
+        path = os.path.join(self.wd, 'a.yaml')
 
-        sys.stdout, sys.stderr = StringIO(), StringIO()
-        with self.assertRaises(SystemExit) as ctx:
-            cli.run((file, '--format', 'auto'))
-
-        self.assertEqual(ctx.exception.code, 1)
-
-        out, err = sys.stdout.getvalue(), sys.stderr.getvalue()
-        self.assertEqual(out, (
+        with RunContext(self) as ctx:
+            cli.run((path, '--format', 'auto'))
+        expected_out = (
             '%s\n'
             '  2:4       error    trailing spaces  (trailing-spaces)\n'
             '  3:4       error    no new line character at the end of file  '
             '(new-line-at-end-of-file)\n'
-            '\n' % file))
-        self.assertEqual(err, '')
+            '\n' % path)
+        self.assertEqual(
+            (ctx.returncode, ctx.stdout, ctx.stderr), (1, expected_out, ''))
 
     def test_run_format_colored(self):
-        file = os.path.join(self.wd, 'a.yaml')
-
-        sys.stdout, sys.stderr = StringIO(), StringIO()
-        with self.assertRaises(SystemExit) as ctx:
-            cli.run((file, '--format', 'colored'))
-
-        self.assertEqual(ctx.exception.code, 1)
+        path = os.path.join(self.wd, 'a.yaml')
 
-        out, err = sys.stdout.getvalue(), sys.stderr.getvalue()
-        self.assertEqual(out, (
+        with RunContext(self) as ctx:
+            cli.run((path, '--format', 'colored'))
+        expected_out = (
             '\033[4m%s\033[0m\n'
             '  \033[2m2:4\033[0m       \033[31merror\033[0m    '
             'trailing spaces  \033[2m(trailing-spaces)\033[0m\n'
             '  \033[2m3:4\033[0m       \033[31merror\033[0m    '
             'no new line character at the end of file  '
             '\033[2m(new-line-at-end-of-file)\033[0m\n'
-            '\n' % file))
-        self.assertEqual(err, '')
+            '\n' % path)
+        self.assertEqual(
+            (ctx.returncode, ctx.stdout, ctx.stderr), (1, expected_out, ''))
 
     def test_run_read_from_stdin(self):
         # prepares stdin with an invalid yaml string so that we can check
         # for its specific error, and be assured that stdin was read
-        sys.stdout, sys.stderr = StringIO(), StringIO()
+        self.addCleanup(setattr, sys, 'stdin', sys.__stdin__)
         sys.stdin = StringIO(
             'I am a string\n'
             'therefore: I am an error\n')
 
-        with self.assertRaises(SystemExit) as ctx:
+        with RunContext(self) as ctx:
             cli.run(('-', '-f', 'parsable'))
-
-        self.assertNotEqual(ctx.exception.code, 0)
-
-        out, err = sys.stdout.getvalue(), sys.stderr.getvalue()
-        self.assertEqual(out, (
+        expected_out = (
             'stdin:2:10: [error] syntax error: '
-            'mapping values are not allowed here (syntax)\n'))
-        self.assertEqual(err, '')
+            'mapping values are not allowed here (syntax)\n')
+        self.assertEqual(
+            (ctx.returncode, ctx.stdout, ctx.stderr), (1, expected_out, ''))
 
     def test_run_no_warnings(self):
-        file = os.path.join(self.wd, 'a.yaml')
-
-        sys.stdout, sys.stderr = StringIO(), StringIO()
-        with self.assertRaises(SystemExit) as ctx:
-            cli.run((file, '--no-warnings', '-f', 'auto'))
-
-        self.assertEqual(ctx.exception.code, 1)
+        path = os.path.join(self.wd, 'a.yaml')
 
-        out, err = sys.stdout.getvalue(), sys.stderr.getvalue()
-        self.assertEqual(out, (
+        with RunContext(self) as ctx:
+            cli.run((path, '--no-warnings', '-f', 'auto'))
+        expected_out = (
             '%s\n'
             '  2:4       error    trailing spaces  (trailing-spaces)\n'
             '  3:4       error    no new line character at the end of file  '
             '(new-line-at-end-of-file)\n'
-            '\n' % file))
-        self.assertEqual(err, '')
-
-        file = os.path.join(self.wd, 'warn.yaml')
+            '\n' % path)
+        self.assertEqual(
+            (ctx.returncode, ctx.stdout, ctx.stderr), (1, expected_out, ''))
 
-        sys.stdout, sys.stderr = StringIO(), StringIO()
-        with self.assertRaises(SystemExit) as ctx:
-            cli.run((file, '--no-warnings', '-f', 'auto'))
+        path = os.path.join(self.wd, 'warn.yaml')
 
-        self.assertEqual(ctx.exception.code, 0)
+        with RunContext(self) as ctx:
+            cli.run((path, '--no-warnings', '-f', 'auto'))
+        self.assertEqual(ctx.returncode, 0)
 
     def test_run_no_warnings_and_strict(self):
-        file = os.path.join(self.wd, 'warn.yaml')
-
-        sys.stdout, sys.stderr = StringIO(), StringIO()
-        with self.assertRaises(SystemExit) as ctx:
-            cli.run((file, '--no-warnings', '-s'))
+        path = os.path.join(self.wd, 'warn.yaml')
 
-        self.assertEqual(ctx.exception.code, 2)
+        with RunContext(self) as ctx:
+            cli.run((path, '--no-warnings', '-s'))
+        self.assertEqual(ctx.returncode, 2)