a useful decorator for cleaning up threads

2 files changed, 13 insertions, 7 deletions

diff --git a/Lib/test/test_hashlib.py b/Lib/test/test_hashlib.py
index b736755..e77b246 100644
--- a/Lib/test/test_hashlib.py
+++ b/Lib/test/test_hashlib.py
@@ -254,14 +254,9 @@ class HashLibTestCase(unittest.TestCase):
 
         self.assertEqual(expected_hash, hasher.hexdigest())
 
-
+@test_support.reap_threads
 def test_main():
-    key = test_support.threading_setup()
-    try:
-        test_support.run_unittest(HashLibTestCase)
-    finally:
-        test_support.threading_cleanup(*key)
-
+    test_support.run_unittest(HashLibTestCase)
 
 if __name__ == "__main__":
     test_main()
diff --git a/Lib/test/test_support.py b/Lib/test/test_support.py
index 7d0ec8d..52dd16e 100644
--- a/Lib/test/test_support.py
+++ b/Lib/test/test_support.py
@@ -5,6 +5,7 @@ if __name__ != 'test.test_support':
 
 import contextlib
 import errno
+import functools
 import socket
 import sys
 import os
@@ -934,6 +935,16 @@ def threading_cleanup(num_active, num_limbo):
         count += 1
         time.sleep(0.1)
 
+def reap_threads(func):
+    @functools.wraps(func)
+    def decorator(*args):
+        key = threading_setup()
+        try:
+            return func(*args)
+        finally:
+            threading_cleanup(*key)
+    return decorator
+
 def reap_children():
     """Use this function at the end of test_main() whenever sub-processes
     are started.  This will help ensure that no extra children (zombies)