Fix test_unpickler_on_missing test

pesiwang · Nov 30, 2023 · 3ec3570 · 3ec3570
1 parent 060c40a
commit 3ec3570
Showing 1 changed file with 34 additions and 32 deletions.
diff --git a/tests/jsonpickle_test.py b/tests/jsonpickle_test.py
@@ -498,42 +498,44 @@ def test_restore_legacy_builtins(self):
         self.assertTrue(cls is int)
 
     def test_unpickler_on_missing(self):
-        class SimpleClass(object):
-            def __init__(self, i):
-                self.i = i
+        encoded = jsonpickle.encode(Outer.Middle.Inner())
+        self.assertTrue(isinstance(jsonpickle.decode(encoded), Outer.Middle.Inner))
+
+        # Alter the encoded string to create cases where the class is missing, in multiple levels
+        self.assertTrue(encoded == '{"py/object": "jsonpickle_test.Outer.Middle.Inner"}')
+        missing_cases = [
+            '{"py/object": "MISSING.Outer.Middle.Inner"}',
+            '{"py/object": "jsonpickle_test.MISSING.Middle.Inner"}',
+            '{"py/object": "jsonpickle_test.Outer.MISSING.Inner"}',
+            '{"py/object": "jsonpickle_test.Outer.Middle.MISSING"}',
+        ]
+
+        for case in missing_cases:
+            # https://docs.python.org/3/library/warnings.html#testing-warnings
+            with warnings.catch_warnings(record=True) as w:
+                warnings.simplefilter("always")
+                jsonpickle.decode(case, on_missing='warn')
+                self.assertTrue(issubclass(w[-1].category, UserWarning))
+                self.assertTrue(
+                    "Unpickler._restore_object could not find" in str(w[-1].message)
+                )
+
+                jsonpickle.decode(case, on_missing=on_missing_callback)
+                self.assertTrue(issubclass(w[-1].category, RuntimeWarning))
+                self.assertTrue("The unpickler couldn't find" in str(w[-1].message))
 
-        frozen = jsonpickle.encode(SimpleClass(4))
-        del SimpleClass
-
-        # https://docs.python.org/3/library/warnings.html#testing-warnings
-
-        with warnings.catch_warnings(record=True) as w:
-            warnings.simplefilter("always")
-            jsonpickle.decode(frozen, on_missing='warn')
-            self.assertTrue(issubclass(w[-1].category, UserWarning))
             self.assertTrue(
-                "Unpickler._restore_object could not find" in str(w[-1].message)
+                jsonpickle.decode(case, on_missing='ignore')
+                == jsonpickle.backend.json.loads(case)
             )
 
-            jsonpickle.decode(frozen, on_missing=on_missing_callback)
-            self.assertTrue(issubclass(w[-1].category, RuntimeWarning))
-            self.assertTrue("The unpickler couldn't find" in str(w[-1].message))
-
-        self.assertTrue(
-            jsonpickle.decode(frozen, on_missing='ignore')
-            == {
-                'py/object': 'jsonpickle_test.PicklingTestCase.test_unpickler_on_missing.<locals>.SimpleClass',
-                'i': 4,
-            }
-        )
-
-        try:
-            jsonpickle.decode(frozen, on_missing='error')
-        except jsonpickle.errors.ClassNotFoundError:
-            # it's supposed to error
-            self.assertTrue(True)
-        else:
-            self.assertTrue(False)
+            try:
+                jsonpickle.decode(case, on_missing='error')
+            except jsonpickle.errors.ClassNotFoundError:
+                # it's supposed to error
+                self.assertTrue(True)
+            else:
+                self.assertTrue(False)
 
     def test_private_slot_members(self):
         obj = jsonpickle.loads(jsonpickle.dumps(MySlots()))