[Python-checkins] r53631 - sandbox/trunk/import_in_py/test_controlled_importlib.py

[brett.cannon]

Author: brett.cannon
Date: Sat Feb  3 23:44:47 2007
New Revision: 53631

Modified:
   sandbox/trunk/import_in_py/test_controlled_importlib.py
Log:
Test whitelisting of frozen modules.  Also use frozen modules to test that
inra-package relative importing does not break past whitelisting.


Modified: sandbox/trunk/import_in_py/test_controlled_importlib.py
==============================================================================
--- sandbox/trunk/import_in_py/test_controlled_importlib.py	(original)
+++ sandbox/trunk/import_in_py/test_controlled_importlib.py	Sat Feb  3 23:44:47 2007
@@ -1,7 +1,7 @@
 import controlled_importlib
 import mock_importlib
 
-from contextlib import contextmanager
+from contextlib import contextmanager, nested
 import StringIO
 import sys
 import unittest
@@ -127,13 +127,25 @@
     try:
         cached = []
         for name in modules:
-            cached.append(sys.modules[name])
-            del sys.modules[name]
+            if name in sys.modules:
+                cached.append(sys.modules[name])
+                del sys.modules[name]
         yield
     finally:
         for name, module in zip(modules, cached):
             sys.modules[name] = module
 
+ at contextmanager
+def temp_setattr(obj, attribute, value):
+    """Temporarily set an attribute on an object."""
+    try:
+        if hasattr(obj, attribute):
+            old_value = getattr(obj, attribute)
+        setattr(obj, attribute, value)
+        yield value
+    finally:
+        setattr(obj, attribute, old_value)
+
 
 class ControlledImportMethodTests(unittest.TestCase):
 
@@ -171,19 +183,27 @@
         module = mock_importlib.MockModule()
         module_name = '.block'
         with mutate_sys_modules(module, module_name):
-            self.failUnlessRaises(ImportError, import_, module_name, {}, {}, 0)
+            self.failUnlessRaises(ImportError, import_, module_name, level=0)
 
     def test_builtin_whitelisting(self):
         whitelist = sys.builtin_module_names[0]
         blacklist = sys.builtin_module_names[1]
         with remove_from_sys_modules(whitelist, blacklist):
             import_ = controlled_importlib.ControlledImport([whitelist])
-            module = import_(whitelist, {}, {}, 0)
+            module = import_(whitelist, level=0)
             self.failUnlessEqual(module.__name__, whitelist)
-            self.failUnlessRaises(ImportError, import_, blacklist, {}, {}, 0)
+            self.failUnlessRaises(ImportError, import_, blacklist, level=0)
 
     def test_frozen_whitelisting(self):
-        pass
+        whitelist = '__phello__'
+        blacklist = ('__hello__', '__phello__.spam')
+        with nested(temp_setattr(sys, 'stdout', StringIO.StringIO()),
+                remove_from_sys_modules(whitelist, *blacklist)):
+            import_ = controlled_importlib.ControlledImport((), [whitelist])
+            module = import_(whitelist, level=0)
+            self.failUnlessEqual(module.__name__, whitelist)
+            for blacklisted in blacklist:
+                self.failUnlessRaises(ImportError, import_, blacklisted, level=0)
 
     def test_extension_whitelisting(self):
         pass
@@ -201,6 +221,20 @@
         # XXX check both modules, packages, and submodules.
         pass
 
+    def test_relative_import(self):
+        # A relative import within a package should not be able to circumvent
+        # whitelisting.
+        whitelist ='__phello__'
+        blacklist = '__phello__.spam'
+        blacklist_module = blacklist.split('.')[1]
+        with nested(temp_setattr(sys, 'stdout', StringIO.StringIO()),
+                remove_from_sys_modules(whitelist, blacklist)):
+            import_ = controlled_importlib.ControlledImport((), [whitelist])
+            pkg = import_(whitelist, level=0)
+            pkg_fromlist = import_('', pkg.__dict__, {}, [blacklist_module], 1)
+            assert pkg_fromlist.__name__ == whitelist
+            self.failUnless(not hasattr(pkg_fromlist, blacklist_module))
+
 
 def test_main():
     test_support.run_unittest(WhitelistTests,