Test custom reload function

thomasballinger · thomasballinger · commit 2589671d154d · 2015-12-23T09:38:20.000-08:00
diff --git a/bpython/curtsiesfrontend/sitefix.py b/bpython/curtsiesfrontend/sitefix.py
@@ -1,5 +1,6 @@
-import sys
 import functools
+import imp
+import sys
 
 from six.moves import builtins
 
@@ -25,6 +26,8 @@ def monkeypatch_quit():
     import importlib
     if hasattr(importlib, 'reload'):
         orig_reload = importlib.reload
+    elif hasattr(imp, 'reload'):
+        orig_reload = imp.reload
 else:
     orig_reload = builtins.reload
 
@@ -34,18 +37,21 @@ def reload(module):
             orig_stdout = sys.stdout
             orig_stderr = sys.stderr
             orig_stdin = sys.stdin
-            orig_reload(sys)
+            r = orig_reload(sys)
             sys.stdout = orig_stdout
             sys.stderr = orig_stderr
             sys.stdin = orig_stdin
+            return r
         else:
-            orig_reload(sys)
+            return orig_reload(module)
 
     functools.update_wrapper(reload, orig_reload)
 
 
 def monkeypatch_reload():
     if py3 and hasattr(importlib, 'reload'):
         importlib.reload = reload
+    elif py3 and hasattr(imp, 'reload'):
+        imp.reload = reload
     else:
         builtins.reload = reload
diff --git a/bpython/test/test_sitefix.py b/bpython/test/test_sitefix.py
@@ -0,0 +1,42 @@
+import sys
+import optparse
+
+from bpython.test import FixLanguageTestCase as TestCase
+
+from bpython.curtsiesfrontend import sitefix
+
+
+class AttrReplaced(object):
+    def __init__(self, module, attr, replacement):
+        self.module = module
+        self.attr = attr
+        self.replacement = replacement
+
+    def __enter__(self):
+        if hasattr(self.module, self.attr):
+            self.orig_value = getattr(self.module, self.attr)
+        setattr(self.module, self.attr, self.replacement)
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if hasattr(self, 'orig_value'):
+            setattr(self.module, self.attr, self.replacement)
+        else:
+            delattr(self.module, self.attr)
+
+
+class TestCurtsiesReevaluateWithImport(TestCase):
+    def test_reload_works(self):
+        orig = optparse.SUPPRESS_HELP
+        with AttrReplaced(optparse, 'SUPPRESS_HELP', 1):
+            #  reload(optparse)
+            sitefix.reload(optparse)
+            self.assertEqual(optparse.SUPPRESS_HELP, orig)
+
+    def test_reload_sys(self):
+        with AttrReplaced(sys, 'stdin', 1):
+            with AttrReplaced(sys, 'a', 2):
+                with AttrReplaced(sys, 'version', 3):
+                    sitefix.reload(sys)
+                    self.assertEqual(sys.a, 2)  # new attrs stick around
+                    self.assertEqual(sys.stdin, 1)  # stdin stays
+                    self.assertNotEqual(sys.version, 3)  # others replaced
