more test_test refactoring

1 files changed, 37 insertions, 39 deletions

diff --git a/src/distutils2/tests/test_test.py b/src/distutils2/tests/test_test.py
index 6467012..2505d6d 100644
--- a/src/distutils2/tests/test_test.py
+++ b/src/distutils2/tests/test_test.py
@@ -2,9 +2,11 @@ import os
 import re
 import sys
 import shutil
+import unittest as ut1
 
 from copy import copy
 from os.path import join
+from operator import getitem, setitem, delitem
 from StringIO import StringIO
 from distutils2.core import Command
 from distutils2.tests.support import unittest, TempdirManager
@@ -26,28 +28,6 @@ AssertionError: horribly
 
 here = os.path.dirname(os.path.abspath(__file__))
 
-def with_mock_ut2_module(func):
-    def wrapper(*args):
-        class MockUTModule(object): pass
-        class MockUTClass(object):
-            def __init__(*_, **__): pass
-            def discover(self): pass
-            def run(self, _): pass
-        ut2 = MockUTModule()
-        ut2.TestLoader = MockUTClass
-        ut2.TextTestRunner = MockUTClass
-        orig_ut2 = sys.modules.get('unittest2', None)
-        try:
-            sys.modules['unittest2'] = ut2
-            args += (ut2,)
-            func(*args)
-        finally:
-            if orig_ut2 is not None:
-                sys.modules['unittest2'] = orig_ut2
-            else:
-                del sys.modules['unittest2']
-    return wrapper
-
 class TestTest(TempdirManager,
                unittest.TestCase):
 
@@ -78,24 +58,33 @@ class TestTest(TempdirManager,
         shutil.copytree(pkg_dir, temp_pkg_dir)
         return temp_pkg_dir
 
-    def safely_replace(self, obj, attr, new_val=None, delete=False):
+    def safely_replace(self, obj, attr, new_val=None, delete=False, dictionary=False):
         """Replace a object's attribute returning to its original state at the
         end of the test run. Creates the attribute if not present before
         (deleting afterwards). When delete=True, makes sure the value is del'd
-        for the test run."""
-        orig_has_attr = hasattr(obj, attr)
-        orig_val = getattr(obj, attr, None)
+        for the test run. If dictionary is set to True, operates of its items
+        rather than attributes."""
+        if dictionary:
+            _setattr, _getattr, _delattr = setitem, getitem, delitem
+            def _hasattr(_dict, value):
+                return value in _dict
+        else:
+            _setattr, _getattr, _delattr, _hasattr = setattr, getattr, delattr, hasattr
+
+        orig_has_attr = _hasattr(obj, attr)
+        if orig_has_attr:
+            orig_val = _getattr(obj, attr) 
 
         if delete is False:
-            setattr(obj, attr, new_val)
-        elif hasattr(obj, attr):
-            delattr(obj, attr)
+            _setattr(obj, attr, new_val)
+        elif orig_has_attr:
+            _delattr(obj, attr)
 
         def do_cleanup():
             if orig_has_attr:
-                setattr(obj, attr, orig_val)
-            elif hasattr(obj, attr):
-                delattr(obj, attr)
+                _setattr(obj, attr, orig_val)
+            elif _hasattr(obj, attr):
+                _delattr(obj, attr)
 
         self.addCleanup(do_cleanup)
 
@@ -108,7 +97,6 @@ class TestTest(TempdirManager,
             def __init__(*_, **__): pass
             def run(_self, suite):
                 record.append("run")
-        import unittest as ut1
         self.safely_replace(ut1, "TextTestRunner", MockTextTestRunner)
 
         dist = Distribution()
@@ -171,21 +159,31 @@ class TestTest(TempdirManager,
         cmd.run()
         self.assertEqual(["runner called"], record)
 
-    @with_mock_ut2_module
-    def test_gets_unittest_discovery(self, mock_ut2):
+    def prepare_mock_ut2(self):
+        class MockUTClass(object):
+            def __init__(*_, **__): pass
+            def discover(self): pass
+            def run(self, _): pass
+        class MockUTModule(object):
+            TestLoader = MockUTClass
+            TextTestRunner = MockUTClass
+        mock_ut2 = MockUTModule()
+        self.safely_replace(sys.modules, "unittest2", mock_ut2, dictionary=True)
+        return mock_ut2
+
+    def test_gets_unittest_discovery(self):
+        mock_ut2 = self.prepare_mock_ut2()
         dist = Distribution()
         cmd = test(dist)
-        import unittest as ut1
         self.safely_replace(ut1.TestLoader, "discover", lambda: None)
         self.assertEqual(cmd.get_ut_with_discovery(), ut1)
 
         del ut1.TestLoader.discover
         self.assertEqual(cmd.get_ut_with_discovery(), mock_ut2)
 
-    @with_mock_ut2_module
-    def test_calls_discover(self, mock_ut2):
-        import unittest as ut1
+    def test_calls_discover(self):
         self.safely_replace(ut1.TestLoader, "discover", delete=True)
+        mock_ut2 = self.prepare_mock_ut2()
         record = []
         mock_ut2.TestLoader.discover = lambda self, path: record.append(path)
         dist = Distribution()