diff --git a/telegram/ext/dispatcher.py b/telegram/ext/dispatcher.py index 73e2fb096e1..1876a17543d 100644 --- a/telegram/ext/dispatcher.py +++ b/telegram/ext/dispatcher.py @@ -433,10 +433,12 @@ def update_persistence(self): """Update :attr:`user_data` and :attr:`chat_data` in :attr:`persistence`. """ if self.persistence: - for chat_id in self.chat_data: - self.persistence.update_chat_data(chat_id, self.chat_data[chat_id]) - for user_id in self.user_data: - self.persistence.update_user_data(user_id, self.user_data[user_id]) + if self.persistence.store_chat_data: + for chat_id in self.chat_data: + self.persistence.update_chat_data(chat_id, self.chat_data[chat_id]) + if self.persistence.store_user_data: + for user_id in self.user_data: + self.persistence.update_user_data(user_id, self.user_data[user_id]) def add_error_handler(self, callback): """Registers an error handler in the Dispatcher. diff --git a/tests/test_persistence.py b/tests/test_persistence.py index 721e5909b79..0820880b8d9 100644 --- a/tests/test_persistence.py +++ b/tests/test_persistence.py @@ -228,6 +228,24 @@ def pickle_persistence(): on_flush=False) +@pytest.fixture(scope='function') +def pickle_persistence_only_chat(): + return PicklePersistence(filename='pickletest', + store_user_data=False, + store_chat_data=True, + singe_file=False, + on_flush=False) + + +@pytest.fixture(scope='function') +def pickle_persistence_only_user(): + return PicklePersistence(filename='pickletest', + store_user_data=True, + store_chat_data=False, + singe_file=False, + on_flush=False) + + @pytest.fixture(scope='function') def bad_pickle_files(): for name in ['pickletest_user_data', 'pickletest_chat_data', 'pickletest_conversations', @@ -541,7 +559,7 @@ def second(update, context): dp.add_handler(h2) dp.process_update(update) - def test_flush_on_stop(self, bot, update, pickle_persistence, good_pickle_files): + def test_flush_on_stop(self, bot, update, pickle_persistence): u = Updater(bot=bot, persistence=pickle_persistence) dp = u.dispatcher u.running = True @@ -559,6 +577,45 @@ def test_flush_on_stop(self, bot, update, pickle_persistence, good_pickle_files) assert pickle_persistence_2.get_user_data()[4242424242]['my_test'] == 'Working!' assert pickle_persistence_2.get_chat_data()[-4242424242]['my_test2'] == 'Working2!' + def test_flush_on_stop_only_chat(self, bot, update, pickle_persistence_only_chat): + os.remove('pickletest_user_data') + os.remove('pickletest_chat_data') + u = Updater(bot=bot, persistence=pickle_persistence_only_chat) + dp = u.dispatcher + u.running = True + dp.user_data[4242424242]['my_test'] = 'Working!' + dp.chat_data[-4242424242]['my_test2'] = 'Working2!' + u.signal_handler(signal.SIGINT, None) + del (dp) + del (u) + del (pickle_persistence_only_chat) + pickle_persistence_2 = PicklePersistence(filename='pickletest', + store_user_data=False, + store_chat_data=True, + singe_file=False, + on_flush=False) + assert pickle_persistence_2.get_user_data() == {} + assert pickle_persistence_2.get_chat_data()[-4242424242]['my_test2'] == 'Working2!' + + def test_flush_on_stop_only_user(self, bot, update, pickle_persistence_only_user): + os.remove('pickletest_chat_data') + u = Updater(bot=bot, persistence=pickle_persistence_only_user) + dp = u.dispatcher + u.running = True + dp.user_data[4242424242]['my_test'] = 'Working!' + dp.chat_data[-4242424242]['my_test2'] = 'Working2!' + u.signal_handler(signal.SIGINT, None) + del (dp) + del (u) + del (pickle_persistence_only_user) + pickle_persistence_2 = PicklePersistence(filename='pickletest', + store_user_data=True, + store_chat_data=False, + singe_file=False, + on_flush=False) + assert pickle_persistence_2.get_user_data()[4242424242]['my_test'] == 'Working!' + assert pickle_persistence_2.get_chat_data()[-4242424242] == {} + def test_with_conversationHandler(self, dp, update, good_pickle_files, pickle_persistence): dp.persistence = pickle_persistence dp.use_context = True