save & restore warnings correctly in unit tests

madacol · Dec 16, 2017 · 13f3e55 · 13f3e55
1 parent 3ed93b6
commit 13f3e55
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/btcrecover/test/test_passwords.py b/btcrecover/test/test_passwords.py
@@ -44,7 +44,8 @@ class NonClosingBase(object):
 def setUpModule():
     global orig_warnings, tstr, tchr, utf8_opt, BytesIO, StringIO, BytesIONonClosing, StringIONonClosing
 
-    orig_warnings = warnings.filters[:]  # the slice notation takes a shallow copy
+    orig_warnings = warnings.catch_warnings()
+    orig_warnings.__enter__()  # save the current warnings settings (it's a context manager)
     # Convert warnings to errors:
     warnings.simplefilter("error")
     # except this from Intel's OpenCL compiler:
@@ -89,7 +90,7 @@ def close(self): pass
 def tearDownModule():
     global tstr
     tstr = None
-    warnings.filters[:] = orig_warnings  # the slice notation replaces the list contents, not the list itself
+    orig_warnings.__exit__(None, None, None)  # restore the original warnings settings
 
 
 WALLET_DIR = os.path.join(os.path.dirname(__file__), "test-wallets")

diff --git a/btcrecover/test/test_seeds.py b/btcrecover/test/test_seeds.py
@@ -40,12 +40,13 @@
 
 def setUpModule():
     global orig_warnings
-    orig_warnings = warnings.filters[:]  # the slice notation takes a shallow copy
+    orig_warnings = warnings.catch_warnings()
+    orig_warnings.__enter__()  # save the current warnings settings (it's a context manager)
     # Convert warnings to errors:
     warnings.simplefilter("error")
 
 def tearDownModule():
-    warnings.filters[:] = orig_warnings  # the slice notation replaces the list contents, not the list itself
+    orig_warnings.__exit__(None, None, None)  # restore the original warnings settings
 
 
 class TestRecoveryFromWallet(unittest.TestCase):