Do not raise exception on flusher stop (#19)

Co-authored-by: Šimon Let <[email protected]>

Author: Nnonexistent

logtail/flusher.py

@@ -1,8 +1,8 @@
 # coding: utf-8
 from __future__ import print_function, unicode_literals
-import sys
-import time
+
 import threading
+import time
 
 from .compat import queue
 
@@ -17,9 +17,10 @@ def __init__(self, upload, pipe, buffer_capacity, flush_interval):
         self.pipe = pipe
         self.buffer_capacity = buffer_capacity
         self.flush_interval = flush_interval
+        self.should_run = True
 
     def run(self):
-        while True:
+        while self.should_run:
             self.step()
 
     def step(self):
@@ -68,7 +69,7 @@ def step(self):
                 print('Failed to send logs to Better Stack after {} retries: {}'.format(len(RETRY_SCHEDULE), response.exception))
 
         if shutdown and self.pipe.empty():
-            sys.exit(0)
+            self.should_run = False
 
 
 def _initial_time_remaining(flush_interval):

tests/test_flusher.py

@@ -1,6 +1,7 @@
 # coding: utf-8
 from __future__ import print_function, unicode_literals
 import mock
+import sys
 import time
 import threading
 import unittest
@@ -120,8 +121,7 @@ def sleep(time):
         self.assertEqual(self.uploader_calls, len(RETRY_SCHEDULE) + 1)
         self.assertEqual(self.sleep_calls, len(RETRY_SCHEDULE))
 
-    @mock.patch('logtail.flusher.sys.exit')
-    def test_shutdown_condition_empties_queue_and_calls_exit(self, mock_exit):
+    def test_shutdown_condition_empties_queue_and_shuts_down(self):
         self.buffer_capacity = 10
         num_items = 5
         first_frame = list(range(self.buffer_capacity))
@@ -141,4 +141,19 @@ def uploader(frame):
 
         fw.step()
         self.assertEqual(self.upload_calls, 1)
-        self.assertEqual(mock_exit.call_count, 1)
+        self.assertFalse(fw.should_run)
+
+    # test relies on overriding excepthook which is available from 3.8+
+    @unittest.skipIf(sys.version_info < (3, 8), "Test skipped because overriding excepthook is only available on Python 3.8+")
+    def test_shutdown_dont_raise_exception_in_thread(self):
+        original_excepthook = threading.excepthook
+        threading.excepthook = mock.Mock()
+
+        _, _, fw = self._setup_worker()
+        fw.parent_thread = mock.MagicMock(is_alive=lambda: False)
+        fw.step()
+
+        self.assertFalse(fw.should_run)
+        self.assertFalse(threading.excepthook.called)
+
+        threading.excepthook = original_excepthook