1 files changed, 58 insertions, 10 deletions

diff --git a/coverage/test_helpers.py b/coverage/test_helpers.py
index 278b0a8a..50cc3298 100644
--- a/coverage/test_helpers.py
+++ b/coverage/test_helpers.py
@@ -1,7 +1,11 @@
+# Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
+# For details: https://bitbucket.org/ned/coveragepy/src/default/NOTICE.txt
+
 """Mixin classes to help make good tests."""
 
 import atexit
 import collections
+import contextlib
 import os
 import random
 import shutil
@@ -39,6 +43,56 @@ class Tee(object):
             return getattr(self._files[0], name)
 
 
+@contextlib.contextmanager
+def change_dir(new_dir):
+    """Change directory, and then change back.
+
+    Use as a context manager, it will give you the new directory, and later
+    restore the old one.
+
+    """
+    old_dir = os.getcwd()
+    os.chdir(new_dir)
+    try:
+        yield os.getcwd()
+    finally:
+        os.chdir(old_dir)
+
+
+@contextlib.contextmanager
+def saved_sys_path():
+    """Save sys.path, and restore it later."""
+    old_syspath = sys.path[:]
+    try:
+        yield
+    finally:
+        sys.path = old_syspath
+
+
+def setup_with_context_manager(testcase, cm):
+    """Use a contextmanager to setUp a test case.
+
+    If you have a context manager you like::
+
+        with ctxmgr(a, b, c) as v:
+            # do something with v
+
+    and you want to have that effect for a test case, call this function from
+    your setUp, and it will start the context manager for your test, and end it
+    when the test is done::
+
+        def setUp(self):
+            self.v = setup_with_context_manager(self, ctxmgr(a, b, c))
+
+        def test_foo(self):
+            # do something with self.v
+
+    """
+    val = cm.__enter__()
+    testcase.addCleanup(cm.__exit__, None, None, None)
+    return val
+
+
 class ModuleAwareMixin(TestCase):
     """A test case mixin that isolates changes to sys.modules."""
 
@@ -46,7 +100,7 @@ class ModuleAwareMixin(TestCase):
         super(ModuleAwareMixin, self).setUp()
 
         # Record sys.modules here so we can restore it in cleanup_modules.
-        self.old_modules = dict(sys.modules)
+        self.old_modules = list(sys.modules)
         self.addCleanup(self.cleanup_modules)
 
     def cleanup_modules(self):
@@ -64,13 +118,7 @@ class SysPathAwareMixin(TestCase):
 
     def setUp(self):
         super(SysPathAwareMixin, self).setUp()
-
-        self.old_syspath = sys.path[:]
-        self.addCleanup(self.cleanup_syspath)
-
-    def cleanup_syspath(self):
-        """Restore the original sys.path."""
-        sys.path = self.old_syspath
+        setup_with_context_manager(self, saved_sys_path())
 
 
 class EnvironmentAwareMixin(TestCase):
@@ -142,7 +190,7 @@ class TempDirMixin(SysPathAwareMixin, ModuleAwareMixin, TestCase):
     """A test case mixin that creates a temp directory and files in it.
 
     Includes SysPathAwareMixin and ModuleAwareMixin, because making and using
-    temp dirs like this will also need that kind of isolation.
+    temp directories like this will also need that kind of isolation.
 
     """
 
@@ -235,7 +283,7 @@ class TempDirMixin(SysPathAwareMixin, ModuleAwareMixin, TestCase):
 
     # We run some tests in temporary directories, because they may need to make
     # files for the tests. But this is expensive, so we can change per-class
-    # whether a temp dir is used or not.  It's easy to forget to set that
+    # whether a temp directory is used or not.  It's easy to forget to set that
     # option properly, so we track information about what the tests did, and
     # then report at the end of the process on test classes that were set
     # wrong.