ut: cleanup the taking-over-io mechanism.

1) Move the functionality from mock.DiskIO to mock.TestCase itself.
2) Use already existing mechanism (tearDownModules()) for cleaning up
after we override functions in widely used modules.

2 files changed, 26 insertions, 32 deletions

diff --git a/tests/mock/__init__.py b/tests/mock/__init__.py
index fb9dd0725..900bef386 100644
--- a/tests/mock/__init__.py
+++ b/tests/mock/__init__.py
@@ -38,7 +38,7 @@ class TestCase(unittest.TestCase):
     def __init__(self, *args, **kwargs):
         unittest.TestCase.__init__(self, *args, **kwargs)
         self.injectedModules = {}
-        
+
     def setupModules(self, a):
         """Mock specified list of modules and store the list so it can be
            properly unloaded during tearDown"""
@@ -52,12 +52,12 @@ class TestCase(unittest.TestCase):
 
     def modifiedModule(self, mname, mod = None):
         """Mark module (and all it's parents) as tainted"""
-        
+
         oldname=""
         for m in mname.split("."):
             self.injectedModules[oldname+m] = mod
             oldname += m + "."
-        self.injectedModules[mname] = mod    
+        self.injectedModules[mname] = mod
 
     def tearDownModules(self):
         """Unload previously Mocked modules"""
@@ -67,6 +67,27 @@ class TestCase(unittest.TestCase):
         for m in sys.modules.keys():
             if m in self.preexistingModules and not m in self.injectedModules:
                 continue
-            
+
             del sys.modules[m]
 
+    def take_over_io(self, disk, target_module):
+        """ Trick target_module into using disk object as the filesystem.
+
+            This is achieved by overriding the module's 'open' binding as well
+            as many global bindings in os.path.
+        """
+        target_module.open = disk.open
+        self.modifiedModule(target_module.__name__)
+
+        import glob
+        self.modifiedModule("glob")
+        glob.glob = disk.glob_glob
+
+        import os
+        self.modifiedModule("os.path")
+        # this is what os.path implementaion points at, reimport:
+        import posixpath
+        self.modifiedModule("posixpath")
+        os.listdir = disk.os_listdir
+        os.path.exists = disk.os_path_exists
+        os.path.isdir = disk.os_path_isdir
diff --git a/tests/mock/disk.py b/tests/mock/disk.py
index fbd61815f..4fe4c7603 100644
--- a/tests/mock/disk.py
+++ b/tests/mock/disk.py
@@ -18,14 +18,7 @@
 
 from StringIO import StringIO
 import fnmatch
-import glob
-import os.path
-
-_orig_glob_glob = glob.glob
-_orig_open = open
-_orig_os_listdir = os.listdir
-_orig_os_path_exists = os.path.exists
-_orig_os_path_isdir = os.path.isdir
+import os
 
 class DiskIO(object):
     """Simple object to simplify mocking of file operations in Mock
@@ -133,23 +126,3 @@ class DiskIO(object):
 
     def os_access(self, path, mode):
         return self.path_exists(path)
-
-    def take_over_module(self, module):
-        """ Trick module into using this object as the filesystem.
-
-            This is achieved by overriding the module's 'open' binding as well
-            as some bindings in os.path.
-        """
-        module.open = self.open
-        module.glob.glob = self.glob_glob
-        module.os.listdir = self.os_listdir
-        module.os.path.exists = self.os_path_exists
-        module.os.path.isdir = self.os_path_isdir
-
-    @staticmethod
-    def restore_module(module):
-        module.open = _orig_open
-        module.glob.glob = _orig_glob_glob
-        module.os.listdir = _orig_os_listdir
-        module.os.path.exists = _orig_os_path_exists
-        module.os.path.isdir = _orig_os_path_isdir