split inf iterable test for imap and imap_unordered

Author: obaltian

Lib/test/_test_multiprocessing.py

@@ -3138,32 +3138,41 @@ def produce_args():
                 p.terminate()
                 p.join()
 
-    def test_imap_and_imap_unordered_buffersize_on_infinite_iterable(self):
+    def test_imap_and_imap_unordered_buffersize_on_empty_iterable(self):
+        for method_name in ("imap", "imap_unordered"):
+            with self.subTest(method=method_name):
+                method = getattr(self.pool, method_name)
+
+                res = method(str, [], buffersize=2)
+
+                self.assertIsNone(next(res, None))
+
+    def test_imap_buffersize_on_infinite_iterable(self):
         if self.TYPE != "threads":
             self.skipTest("test not appropriate for {}".format(self.TYPE))
 
-        for method_name in ("imap", "imap_unordered"):
-            with self.subTest(method=method_name):
-                p = self.Pool(4)
-                method = getattr(p, method_name)
+        p = self.Pool(4)
+        res = p.imap(str, itertools.count(), buffersize=2)
 
-                res = method(str, itertools.count(), buffersize=2)
+        self.assertEqual(next(res, None), "0")
+        self.assertEqual(next(res, None), "1")
+        self.assertEqual(next(res, None), "2")
 
-                self.assertEqual(next(res, None), "0")
-                self.assertEqual(next(res, None), "1")
-                self.assertEqual(next(res, None), "2")
+        p.terminate()
+        p.join()
 
-                p.terminate()
-                p.join()
+    def test_imap_unordered_buffersize_on_infinite_iterable(self):
+        if self.TYPE != "threads":
+            self.skipTest("test not appropriate for {}".format(self.TYPE))
 
-    def test_imap_and_imap_unordered_buffersize_on_empty_iterable(self):
-        for method_name in ("imap", "imap_unordered"):
-            with self.subTest(method=method_name):
-                method = getattr(self.pool, method_name)
+        p = self.Pool(4)
+        res = p.imap(str, itertools.count(), buffersize=2)
+        first_three_results = sorted(next(res, None) for _ in range(3))
 
-                res = method(str, [], buffersize=2)
+        self.assertEqual(first_three_results, ["0", "1", "2"])
 
-                self.assertIsNone(next(res, None))
+        p.terminate()
+        p.join()
 
     def test_make_pool(self):
         expected_error = (RemoteError if self.TYPE == 'manager'