Add support for post-delivery callbacks

Author: imjoehaines

bugsnag/delivery.py

@@ -32,6 +32,23 @@
 __all__ = ('default_headers', 'Delivery')
 
 
+def _noop():
+    pass
+
+
+def _marshall_post_delivery_callback(post_delivery_callback):
+    if not callable(post_delivery_callback):
+        return _noop
+
+    def safe_post_delivery_callback():
+        try:
+            post_delivery_callback()
+        except Exception:
+            pass
+
+    return safe_post_delivery_callback
+
+
 def create_default_delivery():
     if requests is not None:
         return RequestsDelivery()
@@ -81,6 +98,10 @@ def deliver_sessions(self, config, payload: Any, options=None):
             self.deliver(config, payload, options)
 
     def queue_request(self, request: Callable, config, options: Dict):
+        post_delivery_callback = _marshall_post_delivery_callback(
+            options.pop('post_delivery_callback', None)
+        )
+
         if config.asynchronous and options.pop('asynchronous', True):
             # if an exception escapes the thread, our threading.excepthook
             # will catch it and attempt to deliver it
@@ -91,10 +112,15 @@ def safe_request():
                     request()
                 except Exception as e:
                     config.logger.exception('Notifying Bugsnag failed %s', e)
+                finally:
+                    post_delivery_callback()
 
             Thread(target=safe_request).start()
         else:
-            request()
+            try:
+                request()
+            finally:
+                post_delivery_callback()
 
 
 class UrllibDelivery(Delivery):

tests/test_delivery.py

@@ -1,3 +1,4 @@
+import pytest
 import warnings
 import sys
 
@@ -9,7 +10,7 @@
     DEFAULT_SESSIONS_ENDPOINT
 )
 
-from tests.utils import IntegrationTest
+from tests.utils import BrokenDelivery, IntegrationTest, QueueingDelivery
 
 
 class DeliveryTest(IntegrationTest):
@@ -64,3 +65,135 @@ def test_misconfigured_sessions_endpoint_sends_warning(self):
             delivery.deliver_sessions(self.config, '{"apiKey":"aaab"}')
             self.assertEqual(1, len(warn))
             self.assertEqual(0, len(self.server.events_received))
+
+    def test_it_calls_post_delivery_callback(self):
+        callback_was_called = False
+
+        def post_delivery_callback():
+            nonlocal callback_was_called
+            callback_was_called = True
+
+        self.config.delivery.deliver(
+            self.config,
+            '{"legit": 5}',
+            options={'post_delivery_callback': post_delivery_callback}
+        )
+
+        assert callback_was_called
+
+    def test_it_calls_post_delivery_callback_after_success(self):
+        delivery = QueueingDelivery()
+        self.config.configure(delivery=delivery)
+
+        callback_was_called = False
+
+        def post_delivery_callback():
+            nonlocal callback_was_called
+            callback_was_called = True
+
+        delivery.deliver(
+            self.config,
+            '{"legit": 5}',
+            options={'post_delivery_callback': post_delivery_callback}
+        )
+
+        assert not callback_was_called
+
+        delivery.flush_request_queue()
+        assert callback_was_called
+
+    def test_it_calls_post_delivery_callback_on_failure(self):
+        self.config.configure(delivery=BrokenDelivery())
+
+        callback_was_called = False
+
+        def post_delivery_callback():
+            nonlocal callback_was_called
+            callback_was_called = True
+
+        with pytest.raises(Exception):
+            self.config.delivery.deliver(
+                self.config,
+                '{"legit": 6}',
+                options={'post_delivery_callback': post_delivery_callback}
+            )
+
+        assert callback_was_called
+
+    def test_it_calls_post_delivery_callback_for_sessions(self):
+        callback_was_called = False
+
+        def post_delivery_callback():
+            nonlocal callback_was_called
+            callback_was_called = True
+
+        self.config.delivery.deliver_sessions(
+            self.config,
+            '{"legit": 7}',
+            options={'post_delivery_callback': post_delivery_callback}
+        )
+
+        assert callback_was_called
+
+    def test_it_calls_post_delivery_callback_for_sessions_after_success(self):
+        delivery = QueueingDelivery()
+        self.config.configure(delivery=delivery)
+
+        callback_was_called = False
+
+        def post_delivery_callback():
+            nonlocal callback_was_called
+            callback_was_called = True
+
+        delivery.deliver_sessions(
+            self.config,
+            '{"legit": 5}',
+            options={'post_delivery_callback': post_delivery_callback}
+        )
+
+        assert not callback_was_called
+
+        delivery.flush_request_queue()
+        assert callback_was_called
+
+    def test_it_calls_post_delivery_callback_for_sessions_on_failure(self):
+        self.config.configure(delivery=BrokenDelivery())
+
+        callback_was_called = False
+
+        def post_delivery_callback():
+            nonlocal callback_was_called
+            callback_was_called = True
+
+        with pytest.raises(Exception):
+            self.config.delivery.deliver_sessions(
+                self.config,
+                '{"legit": 8}',
+                options={'post_delivery_callback': post_delivery_callback}
+            )
+
+        assert callback_was_called
+
+    def test_it_handles_invalid_post_delivery_callback(self):
+        self.config.delivery.deliver(
+            self.config,
+            '{"legit": 5}',
+            options={'post_delivery_callback': 'not callable :)'}
+        )
+
+    def test_it_handles_errors_in_post_delivery_callback(self):
+        callback_was_called = False
+
+        def post_delivery_callback():
+            nonlocal callback_was_called
+            callback_was_called = True
+
+            raise Exception('oh dear')
+
+        self.config.delivery.deliver(
+            self.config,
+            '{"legit": 5}',
+            options={'post_delivery_callback': post_delivery_callback}
+        )
+
+        assert callback_was_called

tests/utils.py

@@ -7,6 +7,7 @@
 from http.server import SimpleHTTPRequestHandler, HTTPServer
 
 import bugsnag
+from bugsnag.delivery import Delivery
 
 
 try:
@@ -176,3 +177,27 @@ def sent_session_count(self) -> int:
 class ScaryException(Exception):
     class NestedScaryException(Exception):
         pass
+
+
+class QueueingDelivery(Delivery):
+    def __init__(self):
+        self._queue = []
+
+    def deliver(self, config, payload, options):
+        self._queue.append(options['post_delivery_callback'])
+
+    def flush_request_queue(self):
+        for callback in self._queue:
+            callback()
+
+        self._queue.clear()
+
+
+class BrokenDelivery(Delivery):
+    def deliver(self, config, payload, options):
+        def request():
+            raise Exception('broken!')
+
+        options['asynchronous'] = False
+
+        self.queue_request(request, config, options)