diff --git a/Lib/test/pickletester.py b/Lib/test/pickletester.py index bdc7ef62943a28..4609ebc10672d2 100644 --- a/Lib/test/pickletester.py +++ b/Lib/test/pickletester.py @@ -15,6 +15,7 @@ import types import unittest import weakref +from contextlib import contextmanager from textwrap import dedent from http.cookies import SimpleCookie @@ -152,6 +153,17 @@ def __getinitargs__(self): __main__.E = E E.__module__ = "__main__" +@contextmanager +def fake_main_module(): + real_main = sys.modules.get("__main__") + fake_main = types.ModuleType("__main__") + + sys.modules["__main__"] = fake_main + try: + yield + finally: + sys.modules["__main__"] = real_main + # Simple mutable object. class Object: pass @@ -2287,10 +2299,12 @@ def test_nested_lookup_error(self): obj.__module__ = None for proto in protocols: with self.subTest(proto=proto): - with self.assertRaises(pickle.PicklingError) as cm: - self.dumps(obj, proto) + with fake_main_module(): + with self.assertRaises(pickle.PicklingError) as cm: + self.dumps(obj, proto) self.assertEqual(str(cm.exception), - f"Can't pickle {obj!r}: it's not found as __main__.AbstractPickleTests.spam") + f"Can't pickle {obj!r}: " + "it's not found as __main__.AbstractPickleTests.spam") self.assertEqual(str(cm.exception.__context__), "module '__main__' has no attribute 'AbstractPickleTests'") @@ -2304,16 +2318,19 @@ def test_wrong_object_lookup_error(self): with self.assertRaises(pickle.PicklingError) as cm: self.dumps(obj, proto) self.assertEqual(str(cm.exception), - f"Can't pickle {obj!r}: it's not the same object as {__name__}.AbstractPickleTests") + f"Can't pickle {obj!r}: " + f"it's not the same object as {__name__}.AbstractPickleTests") self.assertIsNone(cm.exception.__context__) obj.__module__ = None for proto in protocols: with self.subTest(proto=proto): - with self.assertRaises(pickle.PicklingError) as cm: - self.dumps(obj, proto) + with fake_main_module(): + with self.assertRaises(pickle.PicklingError) as cm: + self.dumps(obj, proto) self.assertEqual(str(cm.exception), - f"Can't pickle {obj!r}: it's not found as __main__.AbstractPickleTests") + f"Can't pickle {obj!r}: " + "it's not found as __main__.AbstractPickleTests") self.assertEqual(str(cm.exception.__context__), "module '__main__' has no attribute 'AbstractPickleTests'")