add tests for FT build

eendebakpt · May 1, 2025 · May 3, 2025 · May 3, 2025 · 76397799f863f3fcbebefac9f74f913756f16d35
commit 76397799f863f3fcbebefac9f74f913756f16d35
diff --git a/Lib/test/test_free_threading/test_itertools_serialize.py b/Lib/test/test_free_threading/test_itertools_serialize.py
@@ -0,0 +1,82 @@
+import unittest
+from threading import Thread, Barrier
+from itertools import serialize
+from test.support import threading_helper
+
+
+threading_helper.requires_working_threading(module=True)
+
+class non_atomic_iterator:
+
+    def __init__(self, it):
+        self.it = iter(it)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        a = next(self.it)
+        b = next(self.it)
+        return a, b
+
+def count():
+    i = 0
+    while True:
+        i += 1
+        yield i
+
+class SerializeThreading(unittest.TestCase):
+
+    @threading_helper.reap_threads
+    def test_serialize(self):
+        number_of_threads = 10
+        number_of_iterations = 10
+        barrier = Barrier(number_of_threads)
+        def work(it):
+            while True:
+                try:
+                    a, b = next(it)
+                    assert a + 1 == b
+                except StopIteration:
+                    break
+
+        data = tuple(range(400))
+        for it in range(number_of_iterations):
+            serialize_iterator = serialize(non_atomic_iterator(data,))
+            worker_threads = []
+            for ii in range(number_of_threads):
+                worker_threads.append(
+                    Thread(target=work, args=[serialize_iterator]))
+
+            with threading_helper.start_threads(worker_threads):
+                pass
+
+            barrier.reset()
+
+    @threading_helper.reap_threads
+    def test_serialize_generator(self):
+        number_of_threads = 5
+        number_of_iterations = 4
+        barrier = Barrier(number_of_threads)
+        def work(it):
+            barrier.wait()
+            for _ in range(1_000):
+                try:
+                    next(it)
+                except StopIteration:
+                    break
+
+        for it in range(number_of_iterations):
+            generator = serialize(count())
+            worker_threads = []
+            for ii in range(number_of_threads):
+                worker_threads.append(
+                    Thread(target=work, args=[generator]))
+
+            with threading_helper.start_threads(worker_threads):
+                pass
+
+            barrier.reset()
+
+if __name__ == "__main__":
+    unittest.main()