Port stdouts decorator from celery/case

Author: tony

tmuxp/testsuite/helpers.py

@@ -1,11 +1,20 @@
 # -*- coding: utf-8 -*-
-"""Helper methods for tmuxp unittests."""
+"""Helper methods for tmuxp unittests.
+
+_CallableContext, WhateverIO, decorator and stdouts are from the case project,
+https://github.com/celery/case, license BSD 3-clause.
+"""
 
 from __future__ import (absolute_import, division, print_function,
                         unicode_literals, with_statement)
 
 import contextlib
+import functools
+import inspect
+import io
 import logging
+import sys
+from contextlib import contextmanager
 from random import randint
 
 from tmuxp import exc
@@ -194,3 +203,136 @@ def bootstrap(self):
         self.TEST_SESSION_NAME = TEST_SESSION_NAME
         self.server = t
         self.session = session
+
+
+StringIO = io.StringIO
+_SIO_write = StringIO.write
+_SIO_init = StringIO.__init__
+
+
+def update_wrapper(wrapper, wrapped, *args, **kwargs):
+    wrapper = functools.update_wrapper(wrapper, wrapped, *args, **kwargs)
+    wrapper.__wrapped__ = wrapped
+    return wrapper
+
+
+def wraps(wrapped,
+          assigned=functools.WRAPPER_ASSIGNMENTS,
+          updated=functools.WRAPPER_UPDATES):
+    return functools.partial(update_wrapper, wrapped=wrapped,
+                             assigned=assigned, updated=updated)
+
+
+class _CallableContext(object):
+
+    def __init__(self, context, cargs, ckwargs, fun):
+        self.context = context
+        self.cargs = cargs
+        self.ckwargs = ckwargs
+        self.fun = fun
+
+    def __call__(self, *args, **kwargs):
+        return self.fun(*args, **kwargs)
+
+    def __enter__(self):
+        self.ctx = self.context(*self.cargs, **self.ckwargs)
+        return self.ctx.__enter__()
+
+    def __exit__(self, *einfo):
+        if self.ctx:
+            return self.ctx.__exit__(*einfo)
+
+
+def decorator(predicate):
+    context = contextmanager(predicate)
+
+    @wraps(predicate)
+    def take_arguments(*pargs, **pkwargs):
+
+        @wraps(predicate)
+        def decorator(cls):
+            if inspect.isclass(cls):
+                orig_setup = cls.setUp
+                orig_teardown = cls.tearDown
+
+                @wraps(cls.setUp)
+                def around_setup(*args, **kwargs):
+                    try:
+                        contexts = args[0].__rb3dc_contexts__
+                    except AttributeError:
+                        contexts = args[0].__rb3dc_contexts__ = []
+                    p = context(*pargs, **pkwargs)
+                    p.__enter__()
+                    contexts.append(p)
+                    return orig_setup(*args, **kwargs)
+                around_setup.__wrapped__ = cls.setUp
+                cls.setUp = around_setup
+
+                @wraps(cls.tearDown)
+                def around_teardown(*args, **kwargs):
+                    try:
+                        contexts = args[0].__rb3dc_contexts__
+                    except AttributeError:
+                        pass
+                    else:
+                        for context in contexts:
+                            context.__exit__(*sys.exc_info())
+                    orig_teardown(*args, **kwargs)
+                around_teardown.__wrapped__ = cls.tearDown
+                cls.tearDown = around_teardown
+
+                return cls
+            else:
+                @wraps(cls)
+                def around_case(self, *args, **kwargs):
+                    with context(*pargs, **pkwargs) as context_args:
+                        context_args = context_args or ()
+                        if not isinstance(context_args, tuple):
+                            context_args = (context_args,)
+                        return cls(*(self,) + args + context_args, **kwargs)
+                return around_case
+
+        if len(pargs) == 1 and callable(pargs[0]):
+            fun, pargs = pargs[0], ()
+            return decorator(fun)
+        return _CallableContext(context, pargs, pkwargs, decorator)
+    assert take_arguments.__wrapped__
+    return take_arguments
+
+
+class WhateverIO(StringIO):
+
+    def __init__(self, v=None, *a, **kw):
+        _SIO_init(self, v.decode() if isinstance(v, bytes) else v, *a, **kw)
+
+    def write(self, data):
+        _SIO_write(self, data.decode() if isinstance(data, bytes) else data)
+
+
+@decorator
+def stdouts():
+    """Override `sys.stdout` and `sys.stderr` with `StringIO`
+    instances.
+    Decorator example::
+        @mock.stdouts
+        def test_foo(self, stdout, stderr):
+            something()
+            self.assertIn('foo', stdout.getvalue())
+    Context example::
+        with mock.stdouts() as (stdout, stderr):
+            something()
+            self.assertIn('foo', stdout.getvalue())
+    """
+    prev_out, prev_err = sys.stdout, sys.stderr
+    prev_rout, prev_rerr = sys.__stdout__, sys.__stderr__
+    mystdout, mystderr = WhateverIO(), WhateverIO()
+    sys.stdout = sys.__stdout__ = mystdout
+    sys.stderr = sys.__stderr__ = mystderr
+
+    try:
+        yield mystdout, mystderr
+    finally:
+        sys.stdout = prev_out
+        sys.stderr = prev_err
+        sys.__stdout__ = prev_rout
+        sys.__stderr__ = prev_rerr

tmuxp/testsuite/util.py

@@ -15,7 +15,7 @@
 
 from tmuxp import exc
 from tmuxp.exc import BeforeLoadScriptError, BeforeLoadScriptNotExists
-from tmuxp.testsuite.helpers import TestCase, TmuxTestCase
+from tmuxp.testsuite.helpers import TestCase, TmuxTestCase, stdouts
 from tmuxp.util import has_required_tmux_version, run_before_script
 
 logger = logging.getLogger(__name__)
@@ -113,10 +113,12 @@ def test_raise_BeforeLoadScriptError_if_retcode(self):
         with self.assertRaises(BeforeLoadScriptError):
             run_before_script(script_file)
 
-    def test_return_stdout_if_ok(self):
+    @stdouts()
+    def test_return_stdout_if_ok(self, stdout, stderr):
         script_file = os.path.join(fixtures_dir, 'script_complete.sh')
 
         run_before_script(script_file)
+        self.assertIn('hello', stdout.getvalue())
 
 
 class BeforeLoadScriptErrorTestCase(TestCase):

tmuxp/testsuite/workspacebuilder.py

@@ -19,7 +19,7 @@
 
 from tmuxp import Window, config, exc
 from tmuxp._compat import text_type
-from tmuxp.testsuite.helpers import TmuxTestCase
+from tmuxp.testsuite.helpers import TmuxTestCase, stdouts
 from tmuxp.workspacebuilder import WorkspaceBuilder
 
 logger = logging.getLogger(__name__)
@@ -883,8 +883,8 @@ class BeforeLoadScript(TmuxTestCase):
       - pane
     """
 
-    def test_throw_error_if_retcode_error(self):
-
+    @stdouts()
+    def test_throw_error_if_retcode_error(self, stdout, stderr):
         sconfig = kaptan.Kaptan(handler='yaml')
         yaml = self.config_script_fails.format(
             fixtures_dir=fixtures_dir,
@@ -909,8 +909,8 @@ def test_throw_error_if_retcode_error(self):
                 msg="Kills session if before_script exits with errcode"
             )
 
-    def test_throw_error_if_file_not_exists(self):
-
+    @stdouts()
+    def test_throw_error_if_file_not_exists(self, stdout, stderr):
         sconfig = kaptan.Kaptan(handler='yaml')
         yaml = self.config_script_not_exists.format(
             fixtures_dir=fixtures_dir,
@@ -941,7 +941,8 @@ def test_throw_error_if_file_not_exists(self):
                 msg="Kills session if before_script doesn't exist"
             )
 
-    def test_true_if_test_passes(self):
+    @stdouts()
+    def test_true_if_test_passes(self, stdout, stderr):
         assert(
             os.path.exists(os.path.join(fixtures_dir, 'script_complete.sh'))
         )
@@ -960,7 +961,8 @@ def test_true_if_test_passes(self):
         with self.temp_session():
             builder.build(session=self.session)
 
-    def test_true_if_test_passes_with_args(self):
+    @stdouts()
+    def test_true_if_test_passes_with_args(self, *args):
         assert(
             os.path.exists(os.path.join(fixtures_dir, 'script_complete.sh'))
         )