patching __import__ of more versions of python

1 files changed, 15 insertions, 9 deletions

diff --git a/blessed/tests/test_core.py b/blessed/tests/test_core.py
index 3403233..e21a167 100644
--- a/blessed/tests/test_core.py
+++ b/blessed/tests/test_core.py
@@ -405,26 +405,29 @@ def test_win32_missing_tty_modules(monkeypatch):
     "Ensure dummy exception is used when io is without UnsupportedOperation."
     @as_subprocess
     def child():
+        OLD_STYLE = False
         try:
-            import builtins as __builtins__
-        except ImportError:
-            import __builtins__
-
-        original_import = __builtins__.__import__
+            original_import = getattr(__builtins__, '__import__')
+            OLD_STYLE = True
+        except AttributeError:
+            original_import = __builtins__['__import__']
 
         tty_modules = ('termios', 'fcntl', 'tty')
 
-        def __import__(name, *args):
+        def __import__(name, *args, **kwargs):
             if name in tty_modules:
                 raise ImportError
-            return original_import(name, *args)
+            return original_import(name, *args, **kwargs)
 
         for module in tty_modules:
             sys.modules.pop(module, None)
 
         warnings.filterwarnings("error", category=UserWarning)
         try:
-            __builtins__.__import__ = __import__
+            if OLD_STYLE:
+                __builtins__.__import__ = __import__
+            else:
+                __builtins__['__import__'] = __import__
             try:
                 import blessed.terminal
                 imp.reload(blessed.terminal)
@@ -441,7 +444,10 @@ def test_win32_missing_tty_modules(monkeypatch):
             assert term.width == 80
 
         finally:
-            __builtins__.__import__ = original_import
+            if OLD_STYLE:
+                setattr(__builtins__, '__import__', original_import)
+            else:
+                __builtins__['__import__'] = original_import
             warnings.resetwarnings()
             import blessed.terminal
             imp.reload(blessed.terminal)