async def test_custom_pickler_unpickler_with_custom_objects( self, bot, pickle_persistence, good_pickle_files): dict_s = self.DictSub("private", "normal", bot) slot_s = self.SlotsSub("new_var", "private_var") regular = self.NormalClass(12) pickle_persistence.set_bot(bot) await pickle_persistence.update_user_data(1232, { "sub_dict": dict_s, "sub_slots": slot_s, "r": regular }) pp = PicklePersistence("pickletest", single_file=False, on_flush=False) pp.set_bot(bot) # Set the bot data = (await pp.get_user_data())[1232] sub_dict = data["sub_dict"] sub_slots = data["sub_slots"] sub_regular = data["r"] assert sub_dict._bot is bot assert sub_dict.normal == dict_s.normal assert sub_dict._private == dict_s._private assert sub_slots.new_var == slot_s.new_var assert sub_slots._private == slot_s._private assert sub_slots._bot is None # We didn't set the bot, so it shouldn't have it here. assert sub_regular.my_var == regular.my_var
async def test_custom_pickler_unpickler_simple(self, pickle_persistence, update, good_pickle_files, bot, recwarn): pickle_persistence.set_bot( bot) # assign the current bot to the persistence data_with_bot = {"current_bot": update.message} await pickle_persistence.update_chat_data( 12345, data_with_bot) # also calls BotPickler.dumps() # Test that regular pickle load fails - err_msg = ( "A load persistent id instruction was encountered,\nbut no persistent_load " "function was specified.") with pytest.raises(pickle.UnpicklingError, match=err_msg): with open("pickletest_chat_data", "rb") as f: pickle.load(f) # Test that our custom unpickler works as intended -- inserts the current bot # We have to create a new instance otherwise unpickling is skipped pp = PicklePersistence("pickletest", single_file=False, on_flush=False) pp.set_bot(bot) # Set the bot assert (await pp.get_chat_data())[12345]["current_bot"].get_bot() is bot # Now test that pickling of unknown bots in TelegramObjects will be replaced by None- assert not len(recwarn) data_with_bot = {} async with make_bot(token=bot.token) as other_bot: data_with_bot["unknown_bot_in_user"] = User(1, "Dev", False, bot=other_bot) await pickle_persistence.update_chat_data(12345, data_with_bot) assert len(recwarn) == 1 assert recwarn[-1].category is PTBUserWarning assert str( recwarn[-1].message).startswith("Unknown bot instance found.") pp = PicklePersistence("pickletest", single_file=False, on_flush=False) pp.set_bot(bot) assert (await pp.get_chat_data())[12345]["unknown_bot_in_user"]._bot is None