Finish up BasePersistence and implementation

Eldinnie · Eldinnie · commit ecf88fe611e1 · 2018-02-21T09:30:06.000+01:00
diff --git a/telegram/ext/basepersistence.py b/telegram/ext/basepersistence.py
@@ -20,33 +20,61 @@
 
 
 class BasePersistence(object):
+    """Example class for adding persistence to your bot.
+    Subclass this object for different implementations of a persistent bot.
+
+    All relevant methods must be overwritten. This means:
+    * If :attr:`store_chat_data` is ``True`` you must overwrite :attr:`get_chat_data` and
+        :attr:`update_chat_data`
+    * If :attr:`store_user_data` is ``True`` you must overwrite :attr:`get_user_data` and
+        :attr:`update_user_data`
+    * If you want to store conversation data with :class:`telegram.ext.CanversationHandler`, you
+        must overwrite :attr:`get_conversations` and :attr:`update_conversations`
+    * :attr:`flush` will be called when the bot is shutdown, and must always be overwritten.
+
+    Attributes:
+        store_user_data (:obj:`bool`): Whether user_data should be saved by this
+            persistence class.
+        store_chat_data (:obj:`bool`): Whether user_data should be saved by this
+            persistence class
+
+    Args:
+        store_user_data (:obj:`bool`): Whether user_data should be saved by this
+            persistence class.
+        store_chat_data (:obj:`bool`): Whether user_data should be saved by this
+            persistence class
+    """
+
     def __init__(self, store_user_data=False, store_chat_data=False):
-        """
-        Args:
-            store_user_data (:obj:`bool`): Whether user_data should be saved by this
-                persistence class.
-            store_chat_data (:obj:`bool`): Whether user_data should be saved by this
-                persistence class
-        """
         self.store_user_data = store_user_data
         self.store_chat_data = store_chat_data
 
     def get_user_data(self):
-        """"
+        """"Will be called by :class:`telegram.ext.Dispatcher` upon creation with a
+        persistence object. It should return the user_data if stored, or an empty
+        ``defaultdict(dict)``.
+
         Returns:
             :obj:'defaultdict`: The restored user data.
         """
         raise NotImplementedError
 
     def get_chat_data(self):
-        """"
+        """"Will be called by :class:`telegram.ext.Dispatcher` upon creation with a
+        persistence object. It should return the chat_data if stored, or an empty
+        ``defaultdict(dict)``.
+
         Returns:
             :obj:'defaultdict`: The restored chat data.
         """
         raise NotImplementedError
 
     def get_conversations(self, name):
-        """"
+        """"Will be called by :class:`telegram.ext.Dispatcher` when a
+        :class:`telegram.ext.CanversationHandler` is added if
+        :attr:`telegram.ext.ConversationHandler.persistent` is ``True``.
+        It should return the conversations for the handler with `name` or an empty ``dict``
+
         Args:
             name (:obj:`str`): The handlers name.
 
@@ -55,29 +83,40 @@ def get_conversations(self, name):
         """
         raise NotImplementedError
 
-    def update_conversation(self, conversations):
-        """"
+    def update_conversations(self, name, conversations):
+        """"Will be called when a :attr:`telegram.ext.ConversationHandler.update_state`
+        is called. this allows the storeage of the new state in the persistence.
+
         Args:
+            name (:obj:`str`): The handlers name.
             conversation (:obj:'dict`): The :attr:`telegram.ext.ConversationHandler.conversations`
                 dict to store.
         """
         raise NotImplementedError
 
     def update_user_data(self, user_data):
-        """"
+        """"Will be calle by the :class:`telegram.ext.Dispatcher` after a handler has
+        handled an update. It will be the :attr:`telegram.ext.Dispatcher.user_data` defaultdict.
+
         Args:
             user_data (:obj:'defaultdict`): The :attr:`telegram.ext.dispatcher.user_data`
                 defaultdict to store.
         """
         raise NotImplementedError
 
     def update_chat_data(self, chat_data):
-        """"
+        """"Will be calle by the :class:`telegram.ext.Dispatcher` after a handler has
+        handled an update. It will be the :attr:`telegram.ext.Dispatcher.chat_data` defaultdict.
+
         Args:
             chat_data (:obj:'defaultdict`): The :attr:`telegram.ext.dispatcher.chat_data`
                 defaultdict to store.
         """
         raise NotImplementedError
 
     def flush(self):
+        """Will be called by :class:`telegram.ext.Updater` upon receiving a stop signal. Gives the
+        persistence a chance to finish up saving or close a database connection gracefully. If this
+        is not of any importance just pass will be sufficient.
+        """
         raise NotImplementedError
diff --git a/telegram/ext/conversationhandler.py b/telegram/ext/conversationhandler.py
@@ -328,4 +328,4 @@ def update_state(self, new_state, key):
             self.conversations[key] = new_state
 
         if self.persistent:
-            self.persistence.update_conversation(self.conversations)
+            self.persistence.update_conversations(self.name, self.conversations)
diff --git a/telegram/ext/updater.py b/telegram/ext/updater.py
@@ -459,6 +459,8 @@ def _join_threads(self):
 
     def signal_handler(self, signum, frame):
         self.is_idle = False
+        if self.persistence:
+            self.persistence.flush()
         if self.running:
             self.logger.info('Received signal {} ({}), stopping...'.format(
                 signum, get_signal_name(signum)))
diff --git a/tests/test_basepersistence.py b/tests/test_basepersistence.py
@@ -66,7 +66,7 @@ def test_creation(self, base_persistence):
         with pytest.raises(NotImplementedError):
             base_persistence.update_user_data(None)
         with pytest.raises(NotImplementedError):
-            base_persistence.update_conversation(None)
+            base_persistence.update_conversations(None)
         with pytest.raises(NotImplementedError):
             base_persistence.flush()
 
@@ -166,14 +166,28 @@ def callback_unknown_user_or_chat(bot, update, user_data, chat_data):
         rec = caplog.records[-2]
         assert rec.msg == 'Saving chat data raised an error'
         assert rec.levelname == 'ERROR'
-
         m.from_user = user2
         m.chat = chat1
         u = Update(1, m)
         dp.process_update(u)
         m.chat = chat2
         u = Update(2, m)
+
+        def save_chat_data(data):
+            if -987654 not in data:
+                pytest.fail()
+
+        def save_user_data(data):
+            if 54321 not in data:
+                pytest.fail()
+
+        base_persistence.update_chat_data = save_chat_data
+        base_persistence.update_user_data = save_user_data
         dp.process_update(u)
 
         assert dp.user_data[54321][1] == 'test7'
         assert dp.chat_data[-987654][2] == 'test8'
+
+    def test_flush(self, updater):
+        with pytest.raises(NotImplementedError):
+            updater.signal_handler(None, None)
