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
