<commit_before>from __future__ import with_statement

from StringIO import StringIO # No need for cStringIO at this time
from contextlib import contextmanager
from functools import wraps
import copy
import getpass
import sys

from fudge import Fake, patched_context, clear_expectations

from fabric.context_managers import settings
from fabric.network import interpret_host_string
from fabric.state import env
import fabric.network

from server import PORT, mapping, users


class FabricTest(object):
    """
    Nose-oriented test runner class that wipes env after every test.
    """
    def setup(self):
        # Clear Fudge mock expectations
        clear_expectations()
        # Copy env for restoration in teardown
        self.previous_env = copy.deepcopy(env)
        # Set up default networking for test server
        env.disable_known_hosts = True
        interpret_host_string('%s@localhost:%s' % (env.local_user, PORT))
        env.password = <PASSWORD>[env.local_user]
        # Command response mocking is easier without having to account for
        # shell wrapping everywhere.
        env.use_shell = False

    def teardown(self):
        env.update(self.previous_env)


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('stdout', 'stderr')
        def func()
            pass
    """
    def mocked_streams_decorator(func):
        @wraps(func)
        def inner_wrapper(*args, **kwargs):
            if 'stdout' in which:
                my_stdout, sys.stdout = sys.stdout, StringIO()
            if 'stderr' in which:
                my_stderr, sys.stderr = sys.stderr, StringIO()
            result = func(*args, **kwargs)<commit_msg>Tweak mock_streams to be slightly friendlier<commit_after>from __future__ import with_statement

from StringIO import StringIO # No need for cStringIO at this time
from contextlib import contextmanager
from functools import wraps
import copy
import getpass
import sys

from fudge import Fake, patched_context, clear_expectations

from fabric.context_managers import settings
from fabric.network import interpret_host_string
from fabric.state import env
import fabric.network

from server import PORT, mapping, users


class FabricTest(object):
    """
    Nose-oriented test runner class that wipes env after every test.
    """
    def setup(self):
        # Clear Fudge mock expectations
        clear_expectations()
        # Copy env for restoration in teardown
        self.previous_env = copy.deepcopy(env)
        # Set up default networking for test server
        env.disable_known_hosts = True
        interpret_host_string('%s@localhost:%s' % (env.local_user, PORT))
        env.password = <PASSWORD>]
        # Command response mocking is easier without having to account for
        # shell wrapping everywhere.
        env.use_shell = False

    def teardown(self):
        env.update(self.previous_env)


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
    """
    which = [which]
    if which == ['both']:
        which = ['stdout', 'stderr']
    def mocked_streams_decorator(func):
        @wraps(func)
        def inner_wrapper(*args, **kwargs):
            if 'stdout' in which:
                my_stdout, sys.stdout = sys.stdout, StringIO()
            if 'stderr' in which:
                my_stderr, sys.stderr = sys.stderr, StringIO()
            result = func(*args, **kwargs)
{'fraction_non_alphanumeric': 0.08325024925224327, 'fraction_numerical': 0.0, 'mean_word_length': 2.651501364877161, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3896765', 'n_tokens_mistral': 1111, 'n_tokens_neox': 1080, 'n_words': 365}