Fix sybrenstuvel#27: Close Pipes at parallel version of genprime.

sybrenstuvel · sybrenstuvel · commit a26f518df5eb · 2016-01-22T11:01:39.000+01:00
Now closing pipes after reading result, and added a unittest.
diff --git a/rsa/parallel.py b/rsa/parallel.py
@@ -65,11 +65,15 @@ def getprime(nbits, poolsize):
     (pipe_recv, pipe_send) = mp.Pipe(duplex=False)
 
     # Create processes
-    procs = [mp.Process(target=_find_prime, args=(nbits, pipe_send))
-             for _ in range(poolsize)]
-    [p.start() for p in procs]
-
-    result = pipe_recv.recv()
+    try:
+        procs = [mp.Process(target=_find_prime, args=(nbits, pipe_send))
+                 for _ in range(poolsize)]
+        [p.start() for p in procs]
+
+        result = pipe_recv.recv()
+    finally:
+        pipe_recv.close()
+        pipe_send.close()
 
     [p.terminate() for p in procs]
 
diff --git a/tests/test_parallel.py b/tests/test_parallel.py
@@ -0,0 +1,20 @@
+"""Test for multiprocess prime generation."""
+
+import unittest
+
+import rsa.prime
+import rsa.parallel
+import rsa.common
+
+
+class ParallelTest(unittest.TestCase):
+    """Tests for multiprocess prime generation."""
+
+    def test_parallel_primegen(self):
+        p = rsa.parallel.getprime(1024, 3)
+
+        self.assertFalse(rsa.prime.is_prime(p - 1))
+        self.assertTrue(rsa.prime.is_prime(p))
+        self.assertFalse(rsa.prime.is_prime(p + 1))
+
+        self.assertEqual(1024, rsa.common.bit_size(p))
