Make sure PP can read files that dont have bot_data (python-telegram-bot#1760)

Bibo-Joshi · web-flow · commit a447760411d5 · 2020-02-08T19:24:35.000+02:00
* Make sure PP can read files that dont have bot_data

* Improve workaround
diff --git a/telegram/ext/picklepersistence.py b/telegram/ext/picklepersistence.py
@@ -84,7 +84,8 @@ def load_singlefile(self):
                 all = pickle.load(f)
                 self.user_data = defaultdict(dict, all['user_data'])
                 self.chat_data = defaultdict(dict, all['chat_data'])
-                self.bot_data = all['bot_data']
+                # For backwards compatibility with files not containing bot data
+                self.bot_data = all.get('bot_data', {})
                 self.conversations = all['conversations']
         except IOError:
             self.conversations = {}
diff --git a/tests/test_persistence.py b/tests/test_persistence.py
@@ -335,6 +335,23 @@ def good_pickle_files(user_data, chat_data, bot_data, conversations):
         os.remove(name)
 
 
+@pytest.fixture(scope='function')
+def pickle_files_wo_bot_data(user_data, chat_data, conversations):
+    all = {'user_data': user_data, 'chat_data': chat_data, 'conversations': conversations}
+    with open('pickletest_user_data', 'wb') as f:
+        pickle.dump(user_data, f)
+    with open('pickletest_chat_data', 'wb') as f:
+        pickle.dump(chat_data, f)
+    with open('pickletest_conversations', 'wb') as f:
+        pickle.dump(conversations, f)
+    with open('pickletest', 'wb') as f:
+        pickle.dump(all, f)
+    yield True
+    for name in ['pickletest_user_data', 'pickletest_chat_data',
+                 'pickletest_conversations', 'pickletest']:
+        os.remove(name)
+
+
 @pytest.fixture(scope='function')
 def update(bot):
     user = User(id=321, first_name='test_user', is_bot=False)
@@ -447,6 +464,54 @@ def test_with_good_single_file(self, pickle_persistence, good_pickle_files):
         with pytest.raises(KeyError):
             conversation2[(123, 123)]
 
+    def test_with_multi_file_wo_bot_data(self, pickle_persistence, pickle_files_wo_bot_data):
+        user_data = pickle_persistence.get_user_data()
+        assert isinstance(user_data, defaultdict)
+        assert user_data[12345]['test1'] == 'test2'
+        assert user_data[67890][3] == 'test4'
+        assert user_data[54321] == {}
+
+        chat_data = pickle_persistence.get_chat_data()
+        assert isinstance(chat_data, defaultdict)
+        assert chat_data[-12345]['test1'] == 'test2'
+        assert chat_data[-67890][3] == 'test4'
+        assert chat_data[-54321] == {}
+
+        bot_data = pickle_persistence.get_bot_data()
+        assert isinstance(bot_data, dict)
+        assert not bot_data.keys()
+
+        conversation1 = 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 = 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)]
+
+    def test_with_single_file_wo_bot_data(self, pickle_persistence, pickle_files_wo_bot_data):
+        pickle_persistence.single_file = True
+        user_data = pickle_persistence.get_user_data()
+        assert isinstance(user_data, defaultdict)
+        assert user_data[12345]['test1'] == 'test2'
+        assert user_data[67890][3] == 'test4'
+        assert user_data[54321] == {}
+
+        chat_data = pickle_persistence.get_chat_data()
+        assert isinstance(chat_data, defaultdict)
+        assert chat_data[-12345]['test1'] == 'test2'
+        assert chat_data[-67890][3] == 'test4'
+        assert chat_data[-54321] == {}
+
+        bot_data = pickle_persistence.get_bot_data()
+        assert isinstance(bot_data, dict)
+        assert not bot_data.keys()
+
     def test_updating_multi_file(self, pickle_persistence, good_pickle_files):
         user_data = pickle_persistence.get_user_data()
         user_data[54321]['test9'] = 'test 10'
@@ -659,7 +724,7 @@ def first(update, context):
             if not context.chat_data == {}:
                 pytest.fail()
             if not context.bot_data == bot_data:
-                pytest.failt()
+                pytest.fail()
             context.user_data['test1'] = 'test2'
             context.chat_data['test3'] = 'test4'
             context.bot_data['test1'] = 'test0'
