refactor tests

Lib/test/pickletester.py

@@ -2408,6 +2408,22 @@ def test_reduce_calls_base(self):
             y = self.loads(s)
             self.assertEqual(y._reduce_called, 1)
 
+    def test_reduce_ex_None(self):
+        c = REX_None()
+        with self.assertRaises(TypeError):
+            self.dumps(c)
+
+    def test_reduce_None(self):
+        c = R_None()
+        with self.assertRaises(TypeError):
+            pickle.dumps(c)
+
+    def test_pickle_setstate_None(self):
+        c = C_None_setstate()
+        p = self.dumps(c)
+        with self.assertRaises(TypeError):
+            self.loads(p)
+
     @no_tracing
     def test_bad_getattr(self):
         # Issue #3514: crash when there is an infinite loop in __getattr__
@@ -3348,6 +3364,20 @@ def __setstate__(self, state):
     def __reduce__(self):
         return type(self), (), self.state
 
+class REX_None:
+    """ Setting __reduce_ex__ to None should fail """
+    __reduce_ex__ = None
+
+class R_None:
+    """ Setting __reduce__ to None should fail """
+    __reduce__ = None
+
+class C_None_setstate:
+    def __getstate__(self):
+        return 1
+
+    __setstate__ = None
+
 
 # Test classes for newobj
 
@@ -3751,6 +3781,25 @@ def test_unpickling_buffering_readline(self):
                 unpickler = self.unpickler_class(f)
                 self.assertEqual(unpickler.load(), data)
 
+    def test_pickle_invalid_reducer_override(self):
+        # gh-103035
+        obj=object()
+
+        f = io.BytesIO()
+        class MyPickler(self.pickler_class):
+            pass
+        pickler = MyPickler(f)
+        pickler.dump(obj)
+
+        pickler.clear_memo()
+        pickler.reducer_override = None
+        with self.assertRaises(TypeError):
+            pickler.dump(obj)
+
+        pickler.clear_memo()
+        pickler.reducer_override = 10
+        with self.assertRaises(TypeError):
+            pickler.dump(obj)
 
 # Tests for dispatch_table attribute
 
@@ -3913,6 +3962,15 @@ def dumps(obj, protocol=None):
 
         self._test_dispatch_table(dumps, dt)
 
+    def test_dispatch_table_None_item(self):
+        # gh-93627
+        obj=object()
+        f = io.BytesIO()
+        pickler = self.pickler_class(f)
+        pickler.dispatch_table = {type(obj): None}
+        with self.assertRaises(TypeError):
+            pickler.dump(obj)
+
     def _test_dispatch_table(self, dumps, dispatch_table):
         def custom_load_dump(obj):
             return pickle.loads(dumps(obj, 0))

Lib/test/test_pickle.py

@@ -570,87 +570,6 @@ def test_multiprocessing_exceptions(self):
                                  ('multiprocessing.context', name))
 
 
-class C_with_reduce_ex:
-    def __reduce_ex__(self, proto):
-        return object.__reduce_ex__(self, proto)
-
-class C_None_reduce_ex:
-    __reduce_ex__ = None
-
-class C_None_reduce:
-    __reduce__ = None
-
-class C_None_setstate:
-    def __getstate__(self):
-        return 1
-
-    __setstate__ = None
-
-class C_setstate:
-    def __init__(self):
-        self.value = 0
-
-    def __getstate__(self):
-        return self.value
-
-    def __setstate__(self, state):
-        print(state)
-        self.value = state + 1
-
-class PickleReductionMethodsTests(unittest.TestCase):
-    # see gh-93627
-
-    def test_pickle_invalid_reduction_method(self):
-        c = C_None_reduce_ex()
-        with self.assertRaises(TypeError):
-            pickle.dumps(c)
-
-        c = C_None_reduce()
-        with self.assertRaises(TypeError):
-            pickle.dumps(c)
-
-    def test_pickle_invalid_reducer_override(self):
-        obj=C_with_reduce_ex()
-
-        bio = io.BytesIO()
-        pickler = pickle._Pickler(bio)
-        pickler.dump(obj)
-
-        pickler.clear_memo()
-        pickler.reducer_override = None
-        with self.assertRaises(TypeError):
-            pickler.dump(obj)
-
-        pickler.clear_memo()
-        pickler.reducer_override = 10
-        with self.assertRaises(TypeError):
-            pickler.dump(obj)
-
-    def test_pickle_invalid_dispatch_table(self):
-        obj=C_with_reduce_ex()
-
-        bio = io.BytesIO()
-        pickler = pickle._Pickler(bio)
-        pickler.dispatch_table = {type(obj): None}
-        with self.assertRaises(TypeError):
-            pickler.dump(obj)
-
-
-class PickleSetstateTests(unittest.TestCase):
-    # See gh-103035
-
-    def test_pickle_setstate(self):
-        c = C_setstate()
-        s = pickle.dumps(c)
-        c2 = pickle.loads(s)
-        self.assertEqual(c2.value, c.value + 1)
-
-        c = C_None_setstate()
-        p = pickle.dumps(c)
-        with self.assertRaises(TypeError):
-            pickle.loads(p)
-
-
 def load_tests(loader, tests, pattern):
     tests.addTest(doctest.DocTestSuite())
     return tests