Added test framework for handling module renames.

Factored the import guard in test_py3kwarn.TestStdlibRemovals into
a context manager, namely test_support.CleanImport.

2 files changed, 78 insertions, 12 deletions

diff --git a/Lib/test/test_py3kwarn.py b/Lib/test/test_py3kwarn.py
index 13c0ec3..2bd2a5e 100644
--- a/Lib/test/test_py3kwarn.py
+++ b/Lib/test/test_py3kwarn.py
@@ -1,7 +1,7 @@
 import unittest
 import sys
-from test.test_support import (catch_warning, TestSkipped, run_unittest,
-                                TestSkipped)
+from test.test_support import (catch_warning, CleanImport,
+                               TestSkipped, run_unittest)
 import warnings
 
 if not sys.py3kwarning:
@@ -138,11 +138,7 @@ class TestStdlibRemovals(unittest.TestCase):
     def check_removal(self, module_name, optional=False):
         """Make sure the specified module, when imported, raises a
         DeprecationWarning and specifies itself in the message."""
-        original_module = None
-        if module_name in sys.modules:
-            original_module = sys.modules[module_name]
-            del sys.modules[module_name]
-        try:
+        with CleanImport(module_name):
             with catch_warning(record=False) as w:
                 warnings.filterwarnings("error", ".+ removed",
                                         DeprecationWarning)
@@ -156,10 +152,6 @@ class TestStdlibRemovals(unittest.TestCase):
                 else:
                     self.fail("DeprecationWarning not raised for %s" %
                                 module_name)
-        finally:
-            if original_module:
-                sys.modules[module_name] = original_module
-
 
     def test_platform_independent_removals(self):
         # Make sure that the modules that are available on all platforms raise
@@ -189,8 +181,49 @@ class TestStdlibRemovals(unittest.TestCase):
             self.assertEquals(str(w.message), msg)
 
 
+class TestStdlibRenames(unittest.TestCase):
+
+    renames = {}
+
+    def check_rename(self, module_name, new_module_name):
+        """Make sure that:
+        - A DeprecationWarning is raised when importing using the
+          old 2.x module name.
+        - The module can be imported using the new 3.x name.
+        - The warning message specify both names.
+        """
+        ModuleType = type(sys) # get the module type object
+        with CleanImport(module_name):
+            with catch_warning(record=False) as w:
+                warnings.filterwarnings("error", ".+ renamed to",
+                                        DeprecationWarning)
+                try:
+                    __import__(module_name, level=0)
+                except DeprecationWarning as exc:
+                    self.assert_(module_name in exc.args[0])
+                    self.assert_(new_module_name in exc.args[0])
+                else:
+                    self.fail("DeprecationWarning not raised for %s" %
+                              module_name)
+        with CleanImport(new_module_name):
+            try:
+                __import__(new_module_name, level=0)
+            except ImportError:
+                self.fail("cannot import %s with its 3.x name, %s" %
+                          module_name, new_module_name)
+            except DeprecationWarning:
+                self.fail("unexpected DeprecationWarning raised for %s" %
+                          module_name)
+
+    def test_module_renames(self):
+        for module_name, new_module_name in self.renames.items():
+            self.check_rename(module_name, new_module_name)
+
+
 def test_main():
-    run_unittest(TestPy3KWarnings, TestStdlibRemovals)
+    run_unittest(TestPy3KWarnings,
+                 TestStdlibRemovals,
+                 TestStdlibRenames)
 
 if __name__ == '__main__':
     test_main()
diff --git a/Lib/test/test_support.py b/Lib/test/test_support.py
index c612c41..d350ed2 100644
--- a/Lib/test/test_support.py
+++ b/Lib/test/test_support.py
@@ -418,6 +418,39 @@ def catch_warning(module=warnings, record=True):
         module.showwarning = original_showwarning
         module.filters = original_filters
 
+
+class CleanImport(object):
+    """Context manager to force import to return a new module reference.
+
+    This is useful for testing module-level behaviours, such as
+    the emission of a DepreciationWarning on import.
+
+    Use like this:
+
+        with CleanImport("foo"):
+            __import__("foo") # new reference
+    """
+
+    def __init__(self, *module_names):
+        self.original_modules = sys.modules.copy()
+        for module_name in module_names:
+            if module_name in sys.modules:
+                module = sys.modules[module_name]
+                # It is possible that module_name is an just alias for
+                # another module (e.g. stub for modules renamed in 3.x).
+                # In that case, we also need delete the real module to clear
+                # the import cache.
+                if module.__name__ != module_name:
+                    del sys.modules[module.__name__]
+                del sys.modules[module_name]
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *ignore_exc):
+        sys.modules.update(self.original_modules)
+
+
 class EnvironmentVarGuard(object):
 
     """Class to help protect the environment variable properly.  Can be used as