python-telegram-bot/tests/test_picklepersistence.py

#!/usr/bin/env python
#
# A library that provides a Python interface to the Telegram Bot API
# Copyright (C) 2015-2022
# Leandro Toledo de Souza <devs@python-telegram-bot.org>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser Public License for more details.
#
# You should have received a copy of the GNU Lesser Public License
# along with this program.  If not, see [http://www.gnu.org/licenses/].
import datetime
import os
import pickle
import gzip
from pathlib import Path

import pytest

from telegram.warnings import PTBUserWarning

from telegram import Update, Message, User, Chat, Bot, TelegramObject
from telegram.ext import (
    PicklePersistence,
    ContextTypes,
    PersistenceInput,
)


@pytest.fixture(autouse=True)
def change_directory(tmp_path: Path):
    orig_dir = Path.cwd()
    # Switch to a temporary directory, so we don't have to worry about cleaning up files
    os.chdir(tmp_path)
    yield
    # Go back to original directory
    os.chdir(orig_dir)


@pytest.fixture(autouse=True)
def reset_callback_data_cache(bot):
    yield
    bot.callback_data_cache.clear_callback_data()
    bot.callback_data_cache.clear_callback_queries()
    bot.arbitrary_callback_data = False


@pytest.fixture(scope="function")
def bot_data():
    return {'test1': 'test2', 'test3': {'test4': 'test5'}}


@pytest.fixture(scope="function")
def chat_data():
    return {-12345: {'test1': 'test2', 'test3': {'test4': 'test5'}}, -67890: {3: 'test4'}}


@pytest.fixture(scope="function")
def user_data():
    return {12345: {'test1': 'test2', 'test3': {'test4': 'test5'}}, 67890: {3: 'test4'}}


@pytest.fixture(scope="function")
def callback_data():
    return [('test1', 1000, {'button1': 'test0', 'button2': 'test1'})], {'test1': 'test2'}


@pytest.fixture(scope='function')
def conversations():
    return {
        'name1': {(123, 123): 3, (456, 654): 4},
        'name2': {(123, 321): 1, (890, 890): 2},
        'name3': {(123, 321): 1, (890, 890): 2},
    }


@pytest.fixture(scope='function')
def pickle_persistence():
    return PicklePersistence(
        filepath='pickletest',
        single_file=False,
        on_flush=False,
    )


@pytest.fixture(scope='function')
def pickle_persistence_only_bot():
    return PicklePersistence(
        filepath='pickletest',
        store_data=PersistenceInput(callback_data=False, user_data=False, chat_data=False),
        single_file=False,
        on_flush=False,
    )


@pytest.fixture(scope='function')
def pickle_persistence_only_chat():
    return PicklePersistence(
        filepath='pickletest',
        store_data=PersistenceInput(callback_data=False, user_data=False, bot_data=False),
        single_file=False,
        on_flush=False,
    )


@pytest.fixture(scope='function')
def pickle_persistence_only_user():
    return PicklePersistence(
        filepath='pickletest',
        store_data=PersistenceInput(callback_data=False, chat_data=False, bot_data=False),
        single_file=False,
        on_flush=False,
    )


@pytest.fixture(scope='function')
def pickle_persistence_only_callback():
    return PicklePersistence(
        filepath='pickletest',
        store_data=PersistenceInput(user_data=False, chat_data=False, bot_data=False),
        single_file=False,
        on_flush=False,
    )


@pytest.fixture(scope='function')
def bad_pickle_files():
    for name in [
        'pickletest_user_data',
        'pickletest_chat_data',
        'pickletest_bot_data',
        'pickletest_callback_data',
        'pickletest_conversations',
        'pickletest',
    ]:
        Path(name).write_text('(())')
    yield True


@pytest.fixture(scope='function')
def invalid_pickle_files():
    for name in [
        'pickletest_user_data',
        'pickletest_chat_data',
        'pickletest_bot_data',
        'pickletest_callback_data',
        'pickletest_conversations',
        'pickletest',
    ]:
        # Just a random way to trigger pickle.UnpicklingError
        # see https://stackoverflow.com/a/44422239/10606962
        with gzip.open(name, 'wb') as file:
            pickle.dump([1, 2, 3], file)
    yield True


@pytest.fixture(scope='function')
def good_pickle_files(user_data, chat_data, bot_data, callback_data, conversations):
    data = {
        'user_data': user_data,
        'chat_data': chat_data,
        'bot_data': bot_data,
        'callback_data': callback_data,
        'conversations': conversations,
    }
    with Path('pickletest_user_data').open('wb') as f:
        pickle.dump(user_data, f)
    with Path('pickletest_chat_data').open('wb') as f:
        pickle.dump(chat_data, f)
    with Path('pickletest_bot_data').open('wb') as f:
        pickle.dump(bot_data, f)
    with Path('pickletest_callback_data').open('wb') as f:
        pickle.dump(callback_data, f)
    with Path('pickletest_conversations').open('wb') as f:
        pickle.dump(conversations, f)
    with Path('pickletest').open('wb') as f:
        pickle.dump(data, f)
    yield True


@pytest.fixture(scope='function')
def pickle_files_wo_bot_data(user_data, chat_data, callback_data, conversations):
    data = {
        'user_data': user_data,
        'chat_data': chat_data,
        'conversations': conversations,
        'callback_data': callback_data,
    }
    with Path('pickletest_user_data').open('wb') as f:
        pickle.dump(user_data, f)
    with Path('pickletest_chat_data').open('wb') as f:
        pickle.dump(chat_data, f)
    with Path('pickletest_callback_data').open('wb') as f:
        pickle.dump(callback_data, f)
    with Path('pickletest_conversations').open('wb') as f:
        pickle.dump(conversations, f)
    with Path('pickletest').open('wb') as f:
        pickle.dump(data, f)
    yield True


@pytest.fixture(scope='function')
def pickle_files_wo_callback_data(user_data, chat_data, bot_data, conversations):
    data = {
        'user_data': user_data,
        'chat_data': chat_data,
        'bot_data': bot_data,
        'conversations': conversations,
    }
    with Path('pickletest_user_data').open('wb') as f:
        pickle.dump(user_data, f)
    with Path('pickletest_chat_data').open('wb') as f:
        pickle.dump(chat_data, f)
    with Path('pickletest_bot_data').open('wb') as f:
        pickle.dump(bot_data, f)
    with Path('pickletest_conversations').open('wb') as f:
        pickle.dump(conversations, f)
    with Path('pickletest').open('wb') as f:
        pickle.dump(data, f)
    yield True


@pytest.fixture(scope='function')
def update(bot):
    user = User(id=321, first_name='test_user', is_bot=False)
    chat = Chat(id=123, type='group')
    message = Message(1, datetime.datetime.now(), chat, from_user=user, text="Hi there", bot=bot)
    return Update(0, message=message)


class TestPicklePersistence:
    """Just tests the PicklePersistence interface. Integration of persistence into Applictation
    is tested in TestBasePersistence!"""

    class DictSub(TelegramObject):  # Used for testing our custom (Un)Pickler.
        def __init__(self, private, normal, b):
            self._private = private
            self.normal = normal
            self._bot = b

    class SlotsSub(TelegramObject):
        __slots__ = ('new_var', '_private')

        def __init__(self, new_var, private):
            self.new_var = new_var
            self._private = private

    class NormalClass:
        def __init__(self, my_var):
            self.my_var = my_var

    @pytest.mark.asyncio
    async def test_slot_behaviour(self, mro_slots, pickle_persistence):
        inst = pickle_persistence
        for attr in inst.__slots__:
            assert getattr(inst, attr, 'err') != 'err', f"got extra slot '{attr}'"
        assert len(mro_slots(inst)) == len(set(mro_slots(inst))), "duplicate slot"

    @pytest.mark.asyncio
    @pytest.mark.parametrize('on_flush', (True, False))
    async def test_on_flush(self, pickle_persistence, on_flush):
        pickle_persistence.on_flush = on_flush
        pickle_persistence.single_file = True
        file_path = Path(pickle_persistence.filepath)

        await pickle_persistence.update_callback_data('somedata')
        assert file_path.is_file() != on_flush

        await pickle_persistence.update_bot_data('data')
        assert file_path.is_file() != on_flush

        await pickle_persistence.update_user_data(123, 'data')
        assert file_path.is_file() != on_flush

        await pickle_persistence.update_chat_data(123, 'data')
        assert file_path.is_file() != on_flush

        await pickle_persistence.update_conversation('name', (1, 1), 'new_state')
        assert file_path.is_file() != on_flush

        await pickle_persistence.flush()
        assert file_path.is_file()

    @pytest.mark.asyncio
    async def test_pickle_behaviour_with_slots(self, pickle_persistence):
        bot_data = await pickle_persistence.get_bot_data()
        bot_data['message'] = Message(3, datetime.datetime.now(), Chat(2, type='supergroup'))
        await pickle_persistence.update_bot_data(bot_data)
        retrieved = await pickle_persistence.get_bot_data()
        assert retrieved == bot_data

    @pytest.mark.asyncio
    async def test_no_files_present_multi_file(self, pickle_persistence):
        assert await pickle_persistence.get_user_data() == {}
        assert await pickle_persistence.get_chat_data() == {}
        assert await pickle_persistence.get_bot_data() == {}
        assert await pickle_persistence.get_callback_data() is None
        assert await pickle_persistence.get_conversations('noname') == {}

    @pytest.mark.asyncio
    async def test_no_files_present_single_file(self, pickle_persistence):
        pickle_persistence.single_file = True
        assert await pickle_persistence.get_user_data() == {}
        assert await pickle_persistence.get_chat_data() == {}
        assert await pickle_persistence.get_bot_data() == {}
        assert await pickle_persistence.get_callback_data() is None
        assert await pickle_persistence.get_conversations('noname') == {}

    @pytest.mark.asyncio
    async def test_with_bad_multi_file(self, pickle_persistence, bad_pickle_files):
        with pytest.raises(TypeError, match='pickletest_user_data'):
            await pickle_persistence.get_user_data()
        with pytest.raises(TypeError, match='pickletest_chat_data'):
            await pickle_persistence.get_chat_data()
        with pytest.raises(TypeError, match='pickletest_bot_data'):
            await pickle_persistence.get_bot_data()
        with pytest.raises(TypeError, match='pickletest_callback_data'):
            await pickle_persistence.get_callback_data()
        with pytest.raises(TypeError, match='pickletest_conversations'):
            await pickle_persistence.get_conversations('name')

    @pytest.mark.asyncio
    async def test_with_invalid_multi_file(self, pickle_persistence, invalid_pickle_files):
        with pytest.raises(TypeError, match='pickletest_user_data does not contain'):
            await pickle_persistence.get_user_data()
        with pytest.raises(TypeError, match='pickletest_chat_data does not contain'):
            await pickle_persistence.get_chat_data()
        with pytest.raises(TypeError, match='pickletest_bot_data does not contain'):
            await pickle_persistence.get_bot_data()
        with pytest.raises(TypeError, match='pickletest_callback_data does not contain'):
            await pickle_persistence.get_callback_data()
        with pytest.raises(TypeError, match='pickletest_conversations does not contain'):
            await pickle_persistence.get_conversations('name')

    @pytest.mark.asyncio
    async def test_with_bad_single_file(self, pickle_persistence, bad_pickle_files):
        pickle_persistence.single_file = True
        with pytest.raises(TypeError, match='pickletest'):
            await pickle_persistence.get_user_data()
        with pytest.raises(TypeError, match='pickletest'):
            await pickle_persistence.get_chat_data()
        with pytest.raises(TypeError, match='pickletest'):
            await pickle_persistence.get_bot_data()
        with pytest.raises(TypeError, match='pickletest'):
            await pickle_persistence.get_callback_data()
        with pytest.raises(TypeError, match='pickletest'):
            await pickle_persistence.get_conversations('name')

    @pytest.mark.asyncio
    async def test_with_invalid_single_file(self, pickle_persistence, invalid_pickle_files):
        pickle_persistence.single_file = True
        with pytest.raises(TypeError, match='pickletest does not contain'):
            await pickle_persistence.get_user_data()
        with pytest.raises(TypeError, match='pickletest does not contain'):
            await pickle_persistence.get_chat_data()
        with pytest.raises(TypeError, match='pickletest does not contain'):
            await pickle_persistence.get_bot_data()
        with pytest.raises(TypeError, match='pickletest does not contain'):
            await pickle_persistence.get_callback_data()
        with pytest.raises(TypeError, match='pickletest does not contain'):
            await pickle_persistence.get_conversations('name')

    @pytest.mark.asyncio
    async def test_with_good_multi_file(self, pickle_persistence, good_pickle_files):
        user_data = await pickle_persistence.get_user_data()
        assert isinstance(user_data, dict)
        assert user_data[12345]['test1'] == 'test2'
        assert user_data[67890][3] == 'test4'

        chat_data = await pickle_persistence.get_chat_data()
        assert isinstance(chat_data, dict)
        assert chat_data[-12345]['test1'] == 'test2'
        assert chat_data[-67890][3] == 'test4'

        bot_data = await pickle_persistence.get_bot_data()
        assert isinstance(bot_data, dict)
        assert bot_data['test1'] == 'test2'
        assert bot_data['test3']['test4'] == 'test5'
        assert 'test0' not in bot_data

        callback_data = await pickle_persistence.get_callback_data()
        assert isinstance(callback_data, tuple)
        assert callback_data[0] == [('test1', 1000, {'button1': 'test0', 'button2': 'test1'})]
        assert callback_data[1] == {'test1': 'test2'}

        conversation1 = await pickle_persistence.get_conversations('name1')
        assert isinstance(conversation1, dict)
        assert conversation1[(123, 123)] == 3
        assert conversation1[(456, 654)] == 4
        with pytest.raises(KeyError):
            conversation1[(890, 890)]
        conversation2 = await pickle_persistence.get_conversations('name2')
        assert isinstance(conversation1, dict)
        assert conversation2[(123, 321)] == 1
        assert conversation2[(890, 890)] == 2
        with pytest.raises(KeyError):
            conversation2[(123, 123)]

    @pytest.mark.asyncio
    async def test_with_good_single_file(self, pickle_persistence, good_pickle_files):
        pickle_persistence.single_file = True
        user_data = await pickle_persistence.get_user_data()
        assert isinstance(user_data, dict)
        assert user_data[12345]['test1'] == 'test2'
        assert user_data[67890][3] == 'test4'

        chat_data = await pickle_persistence.get_chat_data()
        assert isinstance(chat_data, dict)
        assert chat_data[-12345]['test1'] == 'test2'
        assert chat_data[-67890][3] == 'test4'

        bot_data = await pickle_persistence.get_bot_data()
        assert isinstance(bot_data, dict)
        assert bot_data['test1'] == 'test2'
        assert bot_data['test3']['test4'] == 'test5'
        assert 'test0' not in bot_data

        callback_data = await pickle_persistence.get_callback_data()
        assert isinstance(callback_data, tuple)
        assert callback_data[0] == [('test1', 1000, {'button1': 'test0', 'button2': 'test1'})]
        assert callback_data[1] == {'test1': 'test2'}

        conversation1 = await pickle_persistence.get_conversations('name1')
        assert isinstance(conversation1, dict)
        assert conversation1[(123, 123)] == 3
        assert conversation1[(456, 654)] == 4
        with pytest.raises(KeyError):
            conversation1[(890, 890)]
        conversation2 = await pickle_persistence.get_conversations('name2')
        assert isinstance(conversation1, dict)
        assert conversation2[(123, 321)] == 1
        assert conversation2[(890, 890)] == 2
        with pytest.raises(KeyError):
            conversation2[(123, 123)]

    @pytest.mark.asyncio
    async def test_with_multi_file_wo_bot_data(self, pickle_persistence, pickle_files_wo_bot_data):
        user_data = await pickle_persistence.get_user_data()
        assert isinstance(user_data, dict)
        assert user_data[12345]['test1'] == 'test2'
        assert user_data[67890][3] == 'test4'

        chat_data = await pickle_persistence.get_chat_data()
        assert isinstance(chat_data, dict)
        assert chat_data[-12345]['test1'] == 'test2'
        assert chat_data[-67890][3] == 'test4'

        bot_data = await pickle_persistence.get_bot_data()
        assert isinstance(bot_data, dict)
        assert not bot_data.keys()

        callback_data = await pickle_persistence.get_callback_data()
        assert isinstance(callback_data, tuple)
        assert callback_data[0] == [('test1', 1000, {'button1': 'test0', 'button2': 'test1'})]
        assert callback_data[1] == {'test1': 'test2'}

        conversation1 = await pickle_persistence.get_conversations('name1')
        assert isinstance(conversation1, dict)
        assert conversation1[(123, 123)] == 3
        assert conversation1[(456, 654)] == 4
        with pytest.raises(KeyError):
            conversation1[(890, 890)]
        conversation2 = await pickle_persistence.get_conversations('name2')
        assert isinstance(conversation1, dict)
        assert conversation2[(123, 321)] == 1
        assert conversation2[(890, 890)] == 2
        with pytest.raises(KeyError):
            conversation2[(123, 123)]

    @pytest.mark.asyncio
    async def test_with_multi_file_wo_callback_data(
        self, pickle_persistence, pickle_files_wo_callback_data
    ):
        user_data = await pickle_persistence.get_user_data()
        assert isinstance(user_data, dict)
        assert user_data[12345]['test1'] == 'test2'
        assert user_data[67890][3] == 'test4'

        chat_data = await pickle_persistence.get_chat_data()
        assert isinstance(chat_data, dict)
        assert chat_data[-12345]['test1'] == 'test2'
        assert chat_data[-67890][3] == 'test4'

        bot_data = await pickle_persistence.get_bot_data()
        assert isinstance(bot_data, dict)
        assert bot_data['test1'] == 'test2'
        assert bot_data['test3']['test4'] == 'test5'
        assert 'test0' not in bot_data

        callback_data = await pickle_persistence.get_callback_data()
        assert callback_data is None

        conversation1 = await pickle_persistence.get_conversations('name1')
        assert isinstance(conversation1, dict)
        assert conversation1[(123, 123)] == 3
        assert conversation1[(456, 654)] == 4
        with pytest.raises(KeyError):
            conversation1[(890, 890)]
        conversation2 = await pickle_persistence.get_conversations('name2')
        assert isinstance(conversation1, dict)
        assert conversation2[(123, 321)] == 1
        assert conversation2[(890, 890)] == 2
        with pytest.raises(KeyError):
            conversation2[(123, 123)]

    @pytest.mark.asyncio
    async def test_with_single_file_wo_bot_data(
        self, pickle_persistence, pickle_files_wo_bot_data
    ):
        pickle_persistence.single_file = True
        user_data = await pickle_persistence.get_user_data()
        assert isinstance(user_data, dict)
        assert user_data[12345]['test1'] == 'test2'
        assert user_data[67890][3] == 'test4'

        chat_data = await pickle_persistence.get_chat_data()
        assert isinstance(chat_data, dict)
        assert chat_data[-12345]['test1'] == 'test2'
        assert chat_data[-67890][3] == 'test4'

        bot_data = await pickle_persistence.get_bot_data()
        assert isinstance(bot_data, dict)
        assert not bot_data.keys()

        callback_data = await pickle_persistence.get_callback_data()
        assert isinstance(callback_data, tuple)
        assert callback_data[0] == [('test1', 1000, {'button1': 'test0', 'button2': 'test1'})]
        assert callback_data[1] == {'test1': 'test2'}

        conversation1 = await pickle_persistence.get_conversations('name1')
        assert isinstance(conversation1, dict)
        assert conversation1[(123, 123)] == 3
        assert conversation1[(456, 654)] == 4
        with pytest.raises(KeyError):
            conversation1[(890, 890)]
        conversation2 = await pickle_persistence.get_conversations('name2')
        assert isinstance(conversation1, dict)
        assert conversation2[(123, 321)] == 1
        assert conversation2[(890, 890)] == 2
        with pytest.raises(KeyError):
            conversation2[(123, 123)]

    @pytest.mark.asyncio
    async def test_with_single_file_wo_callback_data(
        self, pickle_persistence, pickle_files_wo_callback_data
    ):
        user_data = await pickle_persistence.get_user_data()
        assert isinstance(user_data, dict)
        assert user_data[12345]['test1'] == 'test2'
        assert user_data[67890][3] == 'test4'

        chat_data = await pickle_persistence.get_chat_data()
        assert isinstance(chat_data, dict)
        assert chat_data[-12345]['test1'] == 'test2'
        assert chat_data[-67890][3] == 'test4'

        bot_data = await pickle_persistence.get_bot_data()
        assert isinstance(bot_data, dict)
        assert bot_data['test1'] == 'test2'
        assert bot_data['test3']['test4'] == 'test5'
        assert 'test0' not in bot_data

        callback_data = await pickle_persistence.get_callback_data()
        assert callback_data is None

        conversation1 = await pickle_persistence.get_conversations('name1')
        assert isinstance(conversation1, dict)
        assert conversation1[(123, 123)] == 3
        assert conversation1[(456, 654)] == 4
        with pytest.raises(KeyError):
            conversation1[(890, 890)]
        conversation2 = await pickle_persistence.get_conversations('name2')
        assert isinstance(conversation1, dict)
        assert conversation2[(123, 321)] == 1
        assert conversation2[(890, 890)] == 2
        with pytest.raises(KeyError):
            conversation2[(123, 123)]

    @pytest.mark.asyncio
    async def test_updating_multi_file(self, pickle_persistence, good_pickle_files):
        user_data = await pickle_persistence.get_user_data()
        user_data[12345]['test3']['test4'] = 'test6'
        assert pickle_persistence.user_data != user_data
        await pickle_persistence.update_user_data(12345, user_data[12345])
        assert pickle_persistence.user_data == user_data
        with Path('pickletest_user_data').open('rb') as f:
            user_data_test = dict(pickle.load(f))
        assert user_data_test == user_data
        await pickle_persistence.drop_user_data(67890)
        assert 67890 not in await pickle_persistence.get_user_data()

        chat_data = await pickle_persistence.get_chat_data()
        chat_data[-12345]['test3']['test4'] = 'test6'
        assert pickle_persistence.chat_data != chat_data
        await pickle_persistence.update_chat_data(-12345, chat_data[-12345])
        assert pickle_persistence.chat_data == chat_data
        with Path('pickletest_chat_data').open('rb') as f:
            chat_data_test = dict(pickle.load(f))
        assert chat_data_test == chat_data
        await pickle_persistence.drop_chat_data(-67890)
        assert -67890 not in await pickle_persistence.get_chat_data()

        bot_data = await pickle_persistence.get_bot_data()
        bot_data['test3']['test4'] = 'test6'
        assert pickle_persistence.bot_data != bot_data
        await pickle_persistence.update_bot_data(bot_data)
        assert pickle_persistence.bot_data == bot_data
        with Path('pickletest_bot_data').open('rb') as f:
            bot_data_test = pickle.load(f)
        assert bot_data_test == bot_data

        callback_data = await pickle_persistence.get_callback_data()
        callback_data[1]['test3'] = 'test4'
        assert pickle_persistence.callback_data != callback_data
        await pickle_persistence.update_callback_data(callback_data)
        assert pickle_persistence.callback_data == callback_data
        with Path('pickletest_callback_data').open('rb') as f:
            callback_data_test = pickle.load(f)
        assert callback_data_test == callback_data

        conversation1 = await pickle_persistence.get_conversations('name1')
        conversation1[(123, 123)] = 5
        assert not pickle_persistence.conversations['name1'] == conversation1
        await pickle_persistence.update_conversation('name1', (123, 123), 5)
        assert pickle_persistence.conversations['name1'] == conversation1
        assert await pickle_persistence.get_conversations('name1') == conversation1
        with Path('pickletest_conversations').open('rb') as f:
            conversations_test = dict(pickle.load(f))
        assert conversations_test['name1'] == conversation1

        pickle_persistence.conversations = None
        await pickle_persistence.update_conversation('name1', (123, 123), 5)
        assert pickle_persistence.conversations['name1'] == {(123, 123): 5}
        assert await pickle_persistence.get_conversations('name1') == {(123, 123): 5}

    @pytest.mark.asyncio
    async def test_updating_single_file(self, pickle_persistence, good_pickle_files):
        pickle_persistence.single_file = True

        user_data = await pickle_persistence.get_user_data()
        user_data[12345]['test3']['test4'] = 'test6'
        assert pickle_persistence.user_data != user_data
        await pickle_persistence.update_user_data(12345, user_data[12345])
        assert pickle_persistence.user_data == user_data
        with Path('pickletest').open('rb') as f:
            user_data_test = dict(pickle.load(f))['user_data']
        assert user_data_test == user_data
        await pickle_persistence.drop_user_data(67890)
        assert 67890 not in await pickle_persistence.get_user_data()

        chat_data = await pickle_persistence.get_chat_data()
        chat_data[-12345]['test3']['test4'] = 'test6'
        assert pickle_persistence.chat_data != chat_data
        await pickle_persistence.update_chat_data(-12345, chat_data[-12345])
        assert pickle_persistence.chat_data == chat_data
        with Path('pickletest').open('rb') as f:
            chat_data_test = dict(pickle.load(f))['chat_data']
        assert chat_data_test == chat_data
        await pickle_persistence.drop_chat_data(-67890)
        assert -67890 not in await pickle_persistence.get_chat_data()

        bot_data = await pickle_persistence.get_bot_data()
        bot_data['test3']['test4'] = 'test6'
        assert pickle_persistence.bot_data != bot_data
        await pickle_persistence.update_bot_data(bot_data)
        assert pickle_persistence.bot_data == bot_data
        with Path('pickletest').open('rb') as f:
            bot_data_test = pickle.load(f)['bot_data']
        assert bot_data_test == bot_data

        callback_data = await pickle_persistence.get_callback_data()
        callback_data[1]['test3'] = 'test4'
        assert pickle_persistence.callback_data != callback_data
        await pickle_persistence.update_callback_data(callback_data)
        assert pickle_persistence.callback_data == callback_data
        with Path('pickletest').open('rb') as f:
            callback_data_test = pickle.load(f)['callback_data']
        assert callback_data_test == callback_data

        conversation1 = await pickle_persistence.get_conversations('name1')
        conversation1[(123, 123)] = 5
        assert not pickle_persistence.conversations['name1'] == conversation1
        await pickle_persistence.update_conversation('name1', (123, 123), 5)
        assert pickle_persistence.conversations['name1'] == conversation1
        assert await pickle_persistence.get_conversations('name1') == conversation1
        with Path('pickletest').open('rb') as f:
            conversations_test = dict(pickle.load(f))['conversations']
        assert conversations_test['name1'] == conversation1

        pickle_persistence.conversations = None
        await pickle_persistence.update_conversation('name1', (123, 123), 5)
        assert pickle_persistence.conversations['name1'] == {(123, 123): 5}
        assert await pickle_persistence.get_conversations('name1') == {(123, 123): 5}

    @pytest.mark.asyncio
    async def test_updating_single_file_no_data(self, pickle_persistence):
        pickle_persistence.single_file = True
        assert not any(
            [
                pickle_persistence.user_data,
                pickle_persistence.chat_data,
                pickle_persistence.bot_data,
                pickle_persistence.callback_data,
                pickle_persistence.conversations,
            ]
        )
        await pickle_persistence.flush()
        with pytest.raises(FileNotFoundError, match='pickletest'):
            open('pickletest', 'rb')

    @pytest.mark.asyncio
    async def test_save_on_flush_multi_files(self, pickle_persistence, good_pickle_files):
        # Should run without error
        await pickle_persistence.flush()
        pickle_persistence.on_flush = True

        user_data = await pickle_persistence.get_user_data()
        user_data[54321] = {}
        user_data[54321]['test9'] = 'test 10'
        assert pickle_persistence.user_data != user_data

        await pickle_persistence.update_user_data(54321, user_data[54321])
        assert pickle_persistence.user_data == user_data

        await pickle_persistence.drop_user_data(0)
        assert pickle_persistence.user_data == user_data

        with Path('pickletest_user_data').open('rb') as f:
            user_data_test = dict(pickle.load(f))
        assert user_data_test != user_data

        chat_data = await pickle_persistence.get_chat_data()
        chat_data[54321] = {}
        chat_data[54321]['test9'] = 'test 10'
        assert pickle_persistence.chat_data != chat_data

        await pickle_persistence.update_chat_data(54321, chat_data[54321])
        assert pickle_persistence.chat_data == chat_data

        await pickle_persistence.drop_chat_data(0)
        assert pickle_persistence.user_data == user_data

        with Path('pickletest_chat_data').open('rb') as f:
            chat_data_test = dict(pickle.load(f))
        assert chat_data_test != chat_data

        bot_data = await pickle_persistence.get_bot_data()
        bot_data['test6'] = 'test 7'
        assert pickle_persistence.bot_data != bot_data

        await pickle_persistence.update_bot_data(bot_data)
        assert pickle_persistence.bot_data == bot_data

        with Path('pickletest_bot_data').open('rb') as f:
            bot_data_test = pickle.load(f)
        assert bot_data_test != bot_data

        callback_data = await pickle_persistence.get_callback_data()
        callback_data[1]['test3'] = 'test4'
        assert pickle_persistence.callback_data != callback_data

        await pickle_persistence.update_callback_data(callback_data)
        assert pickle_persistence.callback_data == callback_data

        with Path('pickletest_callback_data').open('rb') as f:
            callback_data_test = pickle.load(f)
        assert callback_data_test != callback_data

        conversation1 = await pickle_persistence.get_conversations('name1')
        conversation1[(123, 123)] = 5
        assert not pickle_persistence.conversations['name1'] == conversation1

        await pickle_persistence.update_conversation('name1', (123, 123), 5)
        assert pickle_persistence.conversations['name1'] == conversation1

        with Path('pickletest_conversations').open('rb') as f:
            conversations_test = dict(pickle.load(f))
        assert not conversations_test['name1'] == conversation1

        await pickle_persistence.flush()
        with Path('pickletest_user_data').open('rb') as f:
            user_data_test = dict(pickle.load(f))
        assert user_data_test == user_data

        with Path('pickletest_chat_data').open('rb') as f:
            chat_data_test = dict(pickle.load(f))
        assert chat_data_test == chat_data

        with Path('pickletest_bot_data').open('rb') as f:
            bot_data_test = pickle.load(f)
        assert bot_data_test == bot_data

        with Path('pickletest_conversations').open('rb') as f:
            conversations_test = dict(pickle.load(f))
        assert conversations_test['name1'] == conversation1

    @pytest.mark.asyncio
    async def test_save_on_flush_single_files(self, pickle_persistence, good_pickle_files):
        # Should run without error
        await pickle_persistence.flush()

        pickle_persistence.on_flush = True
        pickle_persistence.single_file = True

        user_data = await pickle_persistence.get_user_data()
        user_data[54321] = {}
        user_data[54321]['test9'] = 'test 10'
        assert pickle_persistence.user_data != user_data
        await pickle_persistence.update_user_data(54321, user_data[54321])
        assert pickle_persistence.user_data == user_data
        with Path('pickletest').open('rb') as f:
            user_data_test = dict(pickle.load(f))['user_data']
        assert user_data_test != user_data

        chat_data = await pickle_persistence.get_chat_data()
        chat_data[54321] = {}
        chat_data[54321]['test9'] = 'test 10'
        assert pickle_persistence.chat_data != chat_data
        await pickle_persistence.update_chat_data(54321, chat_data[54321])
        assert pickle_persistence.chat_data == chat_data
        with Path('pickletest').open('rb') as f:
            chat_data_test = dict(pickle.load(f))['chat_data']
        assert chat_data_test != chat_data

        bot_data = await pickle_persistence.get_bot_data()
        bot_data['test6'] = 'test 7'
        assert pickle_persistence.bot_data != bot_data
        await pickle_persistence.update_bot_data(bot_data)
        assert pickle_persistence.bot_data == bot_data
        with Path('pickletest').open('rb') as f:
            bot_data_test = pickle.load(f)['bot_data']
        assert bot_data_test != bot_data

        callback_data = await pickle_persistence.get_callback_data()
        callback_data[1]['test3'] = 'test4'
        assert pickle_persistence.callback_data != callback_data
        await pickle_persistence.update_callback_data(callback_data)
        assert pickle_persistence.callback_data == callback_data
        with Path('pickletest').open('rb') as f:
            callback_data_test = pickle.load(f)['callback_data']
        assert callback_data_test != callback_data

        conversation1 = await pickle_persistence.get_conversations('name1')
        conversation1[(123, 123)] = 5
        assert not pickle_persistence.conversations['name1'] == conversation1
        await pickle_persistence.update_conversation('name1', (123, 123), 5)
        assert pickle_persistence.conversations['name1'] == conversation1
        with Path('pickletest').open('rb') as f:
            conversations_test = dict(pickle.load(f))['conversations']
        assert not conversations_test['name1'] == conversation1

        await pickle_persistence.flush()
        with Path('pickletest').open('rb') as f:
            user_data_test = dict(pickle.load(f))['user_data']
        assert user_data_test == user_data

        with Path('pickletest').open('rb') as f:
            chat_data_test = dict(pickle.load(f))['chat_data']
        assert chat_data_test == chat_data

        with Path('pickletest').open('rb') as f:
            bot_data_test = pickle.load(f)['bot_data']
        assert bot_data_test == bot_data

        with Path('pickletest').open('rb') as f:
            conversations_test = dict(pickle.load(f))['conversations']
        assert conversations_test['name1'] == conversation1

    @pytest.mark.asyncio
    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 Bot(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

    @pytest.mark.asyncio
    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

    @pytest.mark.parametrize(
        'filepath',
        ['pickletest', Path('pickletest')],
        ids=['str filepath', 'pathlib.Path filepath'],
    )
    @pytest.mark.asyncio
    async def test_filepath_argument_types(self, filepath):
        pick_persist = PicklePersistence(
            filepath=filepath,
            on_flush=False,
        )
        await pick_persist.update_user_data(1, 1)

        assert (await pick_persist.get_user_data())[1] == 1
        assert Path(filepath).is_file()

    @pytest.mark.parametrize('singlefile', [True, False])
    @pytest.mark.parametrize('ud', [int, float, complex])
    @pytest.mark.parametrize('cd', [int, float, complex])
    @pytest.mark.parametrize('bd', [int, float, complex])
    @pytest.mark.asyncio
    async def test_with_context_types(self, ud, cd, bd, singlefile):
        cc = ContextTypes(user_data=ud, chat_data=cd, bot_data=bd)
        persistence = PicklePersistence('pickletest', single_file=singlefile, context_types=cc)

        assert isinstance(await persistence.get_bot_data(), bd)
        assert await persistence.get_bot_data() == 0

        persistence.user_data = None
        persistence.chat_data = None
        await persistence.drop_user_data(123)
        await persistence.drop_chat_data(123)
        assert isinstance(await persistence.get_user_data(), dict)
        assert isinstance(await persistence.get_chat_data(), dict)
        persistence.user_data = None
        persistence.chat_data = None
        await persistence.update_user_data(1, ud(1))
        await persistence.update_chat_data(1, cd(1))
        await persistence.update_bot_data(bd(1))
        assert (await persistence.get_user_data())[1] == 1
        assert (await persistence.get_chat_data())[1] == 1
        assert await persistence.get_bot_data() == 1

        await persistence.flush()
        persistence = PicklePersistence('pickletest', single_file=singlefile, context_types=cc)
        assert isinstance((await persistence.get_user_data())[1], ud)
        assert (await persistence.get_user_data())[1] == 1
        assert isinstance((await persistence.get_chat_data())[1], cd)
        assert (await persistence.get_chat_data())[1] == 1
        assert isinstance(await persistence.get_bot_data(), bd)
        assert await persistence.get_bot_data() == 1

    @pytest.mark.asyncio
    async def test_no_write_if_data_did_not_change(
        self, pickle_persistence, bot_data, user_data, chat_data, conversations, callback_data
    ):
        pickle_persistence.single_file = True
        pickle_persistence.on_flush = False

        await pickle_persistence.update_bot_data(bot_data)
        await pickle_persistence.update_user_data(12345, user_data[12345])
        await pickle_persistence.update_chat_data(-12345, chat_data[-12345])
        await pickle_persistence.update_conversation('name', (1, 1), 'new_state')
        await pickle_persistence.update_callback_data(callback_data)

        assert pickle_persistence.filepath.is_file()
        pickle_persistence.filepath.unlink()
        assert not pickle_persistence.filepath.is_file()

        await pickle_persistence.update_bot_data(bot_data)
        await pickle_persistence.update_user_data(12345, user_data[12345])
        await pickle_persistence.update_chat_data(-12345, chat_data[-12345])
        await pickle_persistence.update_conversation('name', (1, 1), 'new_state')
        await pickle_persistence.update_callback_data(callback_data)

        assert not pickle_persistence.filepath.is_file()