Be more thorough about cleaning up state between tests

colorama/initialise.py

@@ -6,15 +6,24 @@
 from .ansitowin32 import AnsiToWin32
 
 
-orig_stdout = None
-orig_stderr = None
+def _wipe_internal_state_for_tests():
+    global orig_stdout, orig_stderr
+    orig_stdout = None
+    orig_stderr = None
+
+    global wrapped_stdout, wrapped_stderr
+    wrapped_stdout = None
+    wrapped_stderr = None
+
+    global atexit_done
+    atexit_done = False
 
-wrapped_stdout = None
-wrapped_stderr = None
+    global fixed_windows_console
+    fixed_windows_console = False
 
-atexit_done = False
+    # no-op if it wasn't registered
+    atexit.unregister(reset_all)
 
-fixed_windows_console = False
 
 def reset_all():
     if AnsiToWin32 is not None:    # Issue #74: objects might become None at exit
@@ -102,3 +111,7 @@ def wrap_stream(stream, convert, strip, autoreset, wrap):
         if wrapper.should_wrap():
             stream = wrapper.stream
     return stream
+
+
+# Use this for initial setup as well, to reduce code duplication
+_wipe_internal_state_for_tests()

colorama/tests/initialise_test.py

@@ -8,7 +8,7 @@
     from mock import patch
 
 from ..ansitowin32 import StreamWrapper
-from ..initialise import init
+from ..initialise import init, just_fix_windows_console, _wipe_internal_state_for_tests
 from .utils import osname, replace_by
 
 orig_stdout = sys.stdout
@@ -23,6 +23,7 @@ def setUp(self):
         self.assertNotWrapped()
 
     def tearDown(self):
+        _wipe_internal_state_for_tests()
         sys.stdout = orig_stdout
         sys.stderr = orig_stderr