From cc6039adebfebee22da3a6407bcefaa56ce2f5d3 Mon Sep 17 00:00:00 2001
From: Joe Haines <hello@joehaines.co.uk>
Date: Thu, 28 Mar 2024 10:38:02 +0000
Subject: [PATCH] Add flush method that waits for in-flight requests

---
 bugsnag/client.py    | 30 ++++++++++++++++++
 tests/test_client.py | 73 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 103 insertions(+)

diff --git a/bugsnag/client.py b/bugsnag/client.py
index 072445b4..391d1784 100644
--- a/bugsnag/client.py
+++ b/bugsnag/client.py
@@ -336,6 +336,36 @@ def _leave_breadcrumb_for_event(self, event: Event) -> None:
             BreadcrumbType.ERROR
         )
 
+    def flush(self, timeout_ms: int) -> None:
+        # trigger session delivery as there may be outstanding sessions that
+        # haven't been sent yet
+        self.session_tracker.send_sessions()
+
+        stop_event = threading.Event()
+
+        def block_until_no_requests():
+            while (
+                self._request_tracker.has_in_flight_requests() or
+                self.session_tracker._request_tracker.has_in_flight_requests()
+            ):
+                # wait 10ms before checking for in-flight requests again
+                was_stopped = stop_event.wait(0.01)
+
+                # stop checking and exit if the timeout has been exceeded
+                if was_stopped:
+                    break
+
+        thread = threading.Thread(target=block_until_no_requests)
+        thread.start()
+        thread.join(timeout_ms / 1000)
+
+        if thread.is_alive():
+            # tell the thread to stop checking for in-flight requests as the
+            # timeout has been exceeded
+            stop_event.set()
+
+            raise Exception("flush timed out after %dms" % timeout_ms)
+
 
 class ClientContext:
     def __init__(self, client,
diff --git a/tests/test_client.py b/tests/test_client.py
index a4a4f2c0..bfc63dd8 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -1,9 +1,11 @@
 import os
 import re
 import sys
+import time
 import pytest
 import inspect
 import logging
+import threading
 from datetime import datetime, timedelta, timezone
 from unittest.mock import Mock, ANY
 from tests import fixtures
@@ -1749,6 +1751,77 @@ def test_in_flight_event_request_tracking_in_notify_exc_info_failure(self):
 
         assert not client._request_tracker.has_in_flight_requests()
 
+    def test_flush_returns_immediately_when_no_requests_are_outstanding(self):
+        start_s = time.time()
+        self.client.flush(10)
+        end_s = time.time()
+
+        assert end_s - start_s < 0.01
+
+    def test_flush_raises_if_timeout_is_exceeded(self):
+        delivery = QueueingDelivery()
+        configuration = Configuration()
+        configuration.configure(delivery=delivery, api_key='abc')
+
+        client = Client(configuration)
+        client.notify(Exception('oh dear'))
+
+        with pytest.raises(Exception) as exception:
+            client.flush(10)
+
+            assert str(exception) == 'Exception: flush timed out after 10ms'
+
+    def test_flush_waits_for_outstanding_events_before_returning(self):
+        delivery = QueueingDelivery()
+        configuration = Configuration()
+        configuration.configure(delivery=delivery, api_key='abc')
+
+        client = Client(configuration)
+        client.notify(Exception('oh dear'))
+        client.notify(Exception('oh no'))
+        client.notify(Exception('oh my'))
+
+        def flush_request_queue():
+            time.sleep(0.05)
+            assert client._request_tracker.has_in_flight_requests()
+
+            delivery.flush_request_queue()
+            assert not client._request_tracker.has_in_flight_requests()
+
+        thread = threading.Thread(target=flush_request_queue)
+        thread.start()
+
+        client.flush(100)
+
+        # the thread should have stopped before flush could exit
+        assert not thread.is_alive()
+
+    def test_flush_waits_for_outstanding_sessions_before_returning(self):
+        delivery = QueueingDelivery()
+        configuration = Configuration()
+        configuration.configure(delivery=delivery, api_key='abc')
+
+        client = Client(configuration)
+        client.session_tracker.start_session()
+        client.session_tracker.start_session()
+
+        def flush_request_queue():
+            request_tracker = client.session_tracker._request_tracker
+
+            time.sleep(0.05)
+            assert request_tracker.has_in_flight_requests()
+
+            delivery.flush_request_queue()
+            assert not request_tracker.has_in_flight_requests()
+
+        thread = threading.Thread(target=flush_request_queue)
+        thread.start()
+
+        client.flush(100)
+
+        # the thread should have stopped before flush could exit
+        assert not thread.is_alive()
+
 
 @pytest.mark.parametrize("metadata,type", [
     (1234, 'int'),