From a121ab5ecf491db6602cda159e7f1dc01f5c7a1e Mon Sep 17 00:00:00 2001 From: Hinrich Mahler Date: Sun, 4 Aug 2019 13:42:13 +0200 Subject: [PATCH 1/3] Add checks before persisting on stop --- telegram/ext/dispatcher.py | 10 ++++--- tests/test_persistence.py | 53 +++++++++++++++++++++++++++++++++++++- 2 files changed, 58 insertions(+), 5 deletions(-) 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..2dd625e5bc3 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,8 @@ 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, pickle_persistence_only_chat, + pickle_persistence_only_user, good_pickle_files): u = Updater(bot=bot, persistence=pickle_persistence) dp = u.dispatcher u.running = True @@ -559,6 +578,38 @@ 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!' + 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_chat_data()[-4242424242]['my_test2'] == 'Working2!' + + 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!' + def test_with_conversationHandler(self, dp, update, good_pickle_files, pickle_persistence): dp.persistence = pickle_persistence dp.use_context = True From c2eaa9706d88c95b7afde1987be5bcec10faa236 Mon Sep 17 00:00:00 2001 From: Hinrich Mahler Date: Sun, 4 Aug 2019 14:48:59 +0200 Subject: [PATCH 2/3] Fix tests --- tests/test_persistence.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/tests/test_persistence.py b/tests/test_persistence.py index 2dd625e5bc3..ea13b0d1633 100644 --- a/tests/test_persistence.py +++ b/tests/test_persistence.py @@ -230,7 +230,7 @@ def pickle_persistence(): @pytest.fixture(scope='function') def pickle_persistence_only_chat(): - return PicklePersistence(filename='pickletest', + return PicklePersistence(filename='pickletest_chat', store_user_data=False, store_chat_data=True, singe_file=False, @@ -239,7 +239,7 @@ def pickle_persistence_only_chat(): @pytest.fixture(scope='function') def pickle_persistence_only_user(): - return PicklePersistence(filename='pickletest', + return PicklePersistence(filename='pickletest_user', store_user_data=True, store_chat_data=False, singe_file=False, @@ -559,8 +559,7 @@ def second(update, context): dp.add_handler(h2) dp.process_update(update) - def test_flush_on_stop(self, bot, update, pickle_persistence, pickle_persistence_only_chat, - pickle_persistence_only_user, 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 @@ -578,6 +577,7 @@ def test_flush_on_stop(self, bot, update, pickle_persistence, pickle_persistence 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): u = Updater(bot=bot, persistence=pickle_persistence_only_chat) dp = u.dispatcher u.running = True @@ -587,13 +587,15 @@ def test_flush_on_stop(self, bot, update, pickle_persistence, pickle_persistence del (dp) del (u) del (pickle_persistence_only_chat) - pickle_persistence_2 = PicklePersistence(filename='pickletest', + pickle_persistence_2 = PicklePersistence(filename='pickletest_chat', 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): u = Updater(bot=bot, persistence=pickle_persistence_only_user) dp = u.dispatcher u.running = True @@ -603,12 +605,13 @@ def test_flush_on_stop(self, bot, update, pickle_persistence, pickle_persistence del (dp) del (u) del (pickle_persistence_only_user) - pickle_persistence_2 = PicklePersistence(filename='pickletest', + pickle_persistence_2 = PicklePersistence(filename='pickletest_user', 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 @@ -647,7 +650,7 @@ def teardown_class(cls): try: for name in ['pickletest_user_data', 'pickletest_chat_data', 'pickletest_conversations', - 'pickletest']: + 'pickletest', 'pickletest_chat', 'pickletest_user']: os.remove(name) except Exception: pass From 50322a95c8a7317407d5fabd18c42fe31ab5bd2a Mon Sep 17 00:00:00 2001 From: Hinrich Mahler Date: Mon, 5 Aug 2019 10:01:01 +0200 Subject: [PATCH 3/3] Fix tmp file issue in tests --- tests/test_persistence.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/tests/test_persistence.py b/tests/test_persistence.py index ea13b0d1633..0820880b8d9 100644 --- a/tests/test_persistence.py +++ b/tests/test_persistence.py @@ -230,7 +230,7 @@ def pickle_persistence(): @pytest.fixture(scope='function') def pickle_persistence_only_chat(): - return PicklePersistence(filename='pickletest_chat', + return PicklePersistence(filename='pickletest', store_user_data=False, store_chat_data=True, singe_file=False, @@ -239,7 +239,7 @@ def pickle_persistence_only_chat(): @pytest.fixture(scope='function') def pickle_persistence_only_user(): - return PicklePersistence(filename='pickletest_user', + return PicklePersistence(filename='pickletest', store_user_data=True, store_chat_data=False, singe_file=False, @@ -578,6 +578,8 @@ def test_flush_on_stop(self, bot, update, pickle_persistence): 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 @@ -587,7 +589,7 @@ def test_flush_on_stop_only_chat(self, bot, update, pickle_persistence_only_chat del (dp) del (u) del (pickle_persistence_only_chat) - pickle_persistence_2 = PicklePersistence(filename='pickletest_chat', + pickle_persistence_2 = PicklePersistence(filename='pickletest', store_user_data=False, store_chat_data=True, singe_file=False, @@ -596,6 +598,7 @@ def test_flush_on_stop_only_chat(self, bot, update, pickle_persistence_only_chat 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 @@ -605,7 +608,7 @@ def test_flush_on_stop_only_user(self, bot, update, pickle_persistence_only_user del (dp) del (u) del (pickle_persistence_only_user) - pickle_persistence_2 = PicklePersistence(filename='pickletest_user', + pickle_persistence_2 = PicklePersistence(filename='pickletest', store_user_data=True, store_chat_data=False, singe_file=False, @@ -650,7 +653,7 @@ def teardown_class(cls): try: for name in ['pickletest_user_data', 'pickletest_chat_data', 'pickletest_conversations', - 'pickletest', 'pickletest_chat', 'pickletest_user']: + 'pickletest']: os.remove(name) except Exception: pass