Fix for Promise.done_callback (python-telegram-bot#2544)

starry-shivam · Bibo-Joshi · web-flow · commit 46cdeb495a02 · 2021-06-05T17:08:45.000+02:00
* Don't call done_cb on exceptions

Signed-off-by: starry69 &lt;starry369126@outlook.com&gt;

* improve docs

Co-authored-by: Bibo-Joshi &lt;hinrich.mahler@freenet.de&gt;

* revert black

Co-authored-by: Bibo-Joshi &lt;hinrich.mahler@freenet.de&gt;
diff --git a/telegram/ext/utils/promise.py b/telegram/ext/utils/promise.py
@@ -100,7 +100,7 @@ def run(self) -> None:
 
         finally:
             self.done.set()
-            if self._done_callback:
+            if self._exception is None and self._done_callback:
                 try:
                     self._done_callback(self.result())
                 except Exception as exc:
@@ -136,6 +136,10 @@ def add_done_callback(self, callback: Callable) -> None:
         """
         Callback to be run when :class:`telegram.ext.utils.promise.Promise` becomes done.
 
+        Note:
+            Callback won't be called if :attr:`pooled_function`
+            raises an exception.
+
         Args:
             callback (:obj:`callable`): The callable that will be called when promise is done.
             callback will be called by passing ``Promise.result()`` as only positional argument.
diff --git a/tests/test_promise.py b/tests/test_promise.py
@@ -136,3 +136,17 @@ def done_callback(_):
         )
         assert caplog.records[1].message.startswith("Full traceback:")
         assert promise.result() == "done!"
+
+    def test_done_cb_not_run_on_excp(self):
+        def callback():
+            raise TelegramError('Error')
+
+        def done_callback(_):
+            self.test_flag = True
+
+        promise = Promise(callback, [], {})
+        promise.add_done_callback(done_callback)
+        promise.run()
+        assert isinstance(promise.exception, TelegramError)
+        assert promise.done
+        assert self.test_flag is False
