Refactor and migrate tests from test_pickle to pickletester

Author: Aaron Wieczorek

Lib/test/pickletester.py

@@ -2658,7 +2658,6 @@ def test_misc(self):
             self.assert_is_copy(x, y)
 
         # XXX test __reduce__ protocol?
-
     def test_roundtrip_equality(self):
         expected = self._testdata
         for proto in protocols:
@@ -4261,6 +4260,36 @@ def check_array(arr):
         # 2-D, non-contiguous
         check_array(arr[::2])
 
+    def do_test_concurrent_mutation_in_buffer_callback(self, factory):
+        class R:
+            def __bool__(self):
+                buf.release()
+                return True
+
+        max_proto = getattr(self, "proto", pickle.HIGHEST_PROTOCOL)
+        for proto in range(5, max_proto + 1):
+            obj, sub = factory()
+            buf = pickle.PickleBuffer(obj)
+            buffer_callback = lambda _: R()
+
+            with self.subTest(proto=proto, obj=obj, sub=sub):
+                res = self.dumps(buf, proto, buffer_callback=buffer_callback)
+                self.assertIn(sub, res)
+
+    def test_concurrent_mutation_in_buffer_with_bytearray(self):
+        def factory():
+            s = b"a" * 16
+            return bytearray(s), s
+        self.do_test_concurrent_mutation_in_buffer_callback(factory)
+
+    def test_concurrent_mutation_in_buffer_with_memoryview(self):
+        def factory():
+            c, n = b"a", 64
+            sub = c * (n // 2)
+            obj = memoryview(bytearray(c * n))[n // 4 : 3 * n // 4]
+            return obj, sub
+        self.do_test_concurrent_mutation_in_buffer_callback(factory)
+
     def test_evil_class_mutating_dict(self):
         # https://github.com/python/cpython/issues/92930
         from random import getrandbits

Lib/test/test_pickle.py

@@ -388,38 +388,6 @@ class CPicklingErrorTests(PyPicklingErrorTests):
     class CPicklerTests(PyPicklerTests):
         pickler = _pickle.Pickler
         unpickler = _pickle.Unpickler
-        def test_release_in_callback_keepalive(self):
-                base = bytearray(b'A' * 16)
-                pb = pickle.PickleBuffer(base)
-
-                class Evil:
-                    def __bool__(self):
-                        pb.release()
-                        return True
-
-                def callback(p):
-                    return Evil()
-
-                for proto in range(5, pickle.HIGHEST_PROTOCOL + 1):
-                    result = self.dumps(pb, proto, buffer_callback=callback)
-                    self.assertIn(b'A' * 16, result)
-
-        def test_release_in_callback_complex_keepalive(self):
-                base = bytearray(b'A' * 32)
-                view = memoryview(base)[10:26]
-                pb = pickle.PickleBuffer(view)
-
-                class Evil:
-                    def __bool__(self):
-                        pb.release()
-                        return True
-
-                def callback(p):
-                    return Evil()
-
-                for proto in range(5, pickle.HIGHEST_PROTOCOL + 1):
-                    result = self.dumps(pb, proto, buffer_callback=callback)
-                    self.assertIn(b'A' * 16, result)
 
     class CPersPicklerTests(PyPersPicklerTests):
         pickler = _pickle.Pickler