File: mock_streams.py

package info (click to toggle)
fabric 1.14.0-1
  • links: PTS, VCS
  • area: main
  • in suites: bullseye, buster, sid
  • size: 1,240 kB
  • sloc: python: 7,363; makefile: 10
file content (85 lines) | stat: -rw-r--r-- 2,480 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
"""
Stand-alone stream mocking decorator for easier imports.
"""
from functools import wraps
import sys
from StringIO import StringIO  # No need for cStringIO at this time


class CarbonCopy(StringIO):
    """
    A StringIO capable of multiplexing its writes to other buffer objects.
    """

    def __init__(self, buffer='', cc=None):
        """
        If ``cc`` is given and is a file-like object or an iterable of same,
        it/they will be written to whenever this StringIO instance is written
        to.
        """
        StringIO.__init__(self, buffer)
        if cc is None:
            cc = []
        elif hasattr(cc, 'write'):
            cc = [cc]
        self.cc = cc

    def write(self, s):
        StringIO.write(self, s)
        for writer in self.cc:
            writer.write(s)


def mock_streams(which):
    """
    Replaces a stream with a ``StringIO`` during the test, then restores after.

    Must specify which stream (stdout, stderr, etc) via string args, e.g.::

        @mock_streams('stdout')
        def func():
            pass

        @mock_streams('stderr')
        def func():
            pass

        @mock_streams('both')
        def func()
            pass

    If ``'both'`` is specified, not only will both streams be replaced with
    StringIOs, but a new combined-streams output (another StringIO) will appear
    at ``sys.stdall``. This StringIO will resemble what a user sees at a
    terminal, i.e. both streams intermingled.
    """
    both = (which == 'both')
    stdout = (which == 'stdout') or both
    stderr = (which == 'stderr') or both

    def mocked_streams_decorator(func):
        @wraps(func)
        def inner_wrapper(*args, **kwargs):
            if both:
                sys.stdall = StringIO()
                fake_stdout = CarbonCopy(cc=sys.stdall)
                fake_stderr = CarbonCopy(cc=sys.stdall)
            else:
                fake_stdout, fake_stderr = StringIO(), StringIO()
            if stdout:
                my_stdout, sys.stdout = sys.stdout, fake_stdout
            if stderr:
                my_stderr, sys.stderr = sys.stderr, fake_stderr
            try:
                func(*args, **kwargs)
            finally:
                if stdout:
                    sys.stdout = my_stdout
                if stderr:
                    sys.stderr = my_stderr
                if both:
                    del sys.stdall
        return inner_wrapper
    return mocked_streams_decorator