T-9413 Add configurable timeout (#34)

Author: PetrHeinz

logtail/__init__.py

@@ -5,6 +5,6 @@
 from .helpers import LogtailContext, DEFAULT_CONTEXT
 from .formatter import LogtailFormatter
 
-__version__ = '0.3.3'
+__version__ = '0.3.4'
 
 context = DEFAULT_CONTEXT

logtail/handler.py

@@ -16,6 +16,7 @@
 DEFAULT_RAISE_EXCEPTIONS = False
 DEFAULT_DROP_EXTRA_EVENTS = True
 DEFAULT_INCLUDE_EXTRA_ATTRIBUTES = True
+DEFAULT_TIMEOUT = 30
 
 
 class LogtailHandler(logging.Handler):
@@ -29,6 +30,7 @@ def __init__(self,
                  drop_extra_events=DEFAULT_DROP_EXTRA_EVENTS,
                  include_extra_attributes=DEFAULT_INCLUDE_EXTRA_ATTRIBUTES,
                  context=DEFAULT_CONTEXT,
+                 timeout=DEFAULT_TIMEOUT,
                  level=logging.NOTSET):
         super(LogtailHandler, self).__init__(level=level)
         self.source_token = source_token
@@ -38,7 +40,7 @@ def __init__(self,
             self.host = "https://" + host
         self.context = context
         self.pipe = queue.Queue(maxsize=buffer_capacity)
-        self.uploader = Uploader(self.source_token, self.host)
+        self.uploader = Uploader(self.source_token, self.host, timeout)
         self.drop_extra_events = drop_extra_events
         self.include_extra_attributes = include_extra_attributes
         self.buffer_capacity = buffer_capacity

logtail/uploader.py

@@ -9,9 +9,10 @@ def __init__(self, exception):
         self.exception = exception
 
 class Uploader(object):
-    def __init__(self, source_token, host):
+    def __init__(self, source_token, host, timeout):
         self.source_token = source_token
         self.host = host
+        self.timeout = timeout
         self.session = requests.Session()
         self.headers = {
             'Authorization': 'Bearer %s' % source_token,
@@ -21,6 +22,6 @@ def __init__(self, source_token, host):
     def __call__(self, frame):
         data = msgpack.packb(frame, use_bin_type=True)
         try:
-            return self.session.post(self.host, data=data, headers=self.headers)
+            return self.session.post(self.host, data=data, headers=self.headers, timeout=self.timeout)
         except requests.RequestException as e:
-            return Fake500(e)
+            return Fake500(e)

setup.py

@@ -3,7 +3,7 @@
 from setuptools import setup
 
 
-VERSION = '0.3.3'
+VERSION = '0.3.4'
 ROOT_DIR = os.path.dirname(__file__)
 
 REQUIREMENTS = [

tests/test_flusher.py

@@ -20,10 +20,11 @@ class TestFlushWorker(unittest.TestCase):
     buffer_capacity = 5
     flush_interval = 2
     check_interval = 0.01
+    timeout = 0.1
 
     def _setup_worker(self, uploader=None):
         pipe = queue.Queue(maxsize=self.buffer_capacity)
-        uploader = uploader or Uploader(self.source_token, self.host)
+        uploader = uploader or Uploader(self.source_token, self.host, self.timeout)
         fw = FlushWorker(uploader, pipe, self.buffer_capacity, self.flush_interval, self.check_interval)
         return pipe, uploader, fw
 

tests/test_handler.py

@@ -20,6 +20,16 @@ def test_handler_creates_uploader_from_args(self, MockWorker):
         handler = LogtailHandler(source_token=self.source_token, host=self.host)
         self.assertEqual(handler.uploader.source_token, self.source_token)
         self.assertEqual(handler.uploader.host, "https://" + self.host)
+        
+    @patch('logtail.handler.FlushWorker')
+    def test_handler_passes_timeout_to_uploader(self, MockWorker):
+        # Test default timeout
+        handler = LogtailHandler(source_token=self.source_token, host=self.host)
+        self.assertEqual(handler.uploader.timeout, 30)
+        
+        # Test custom timeout
+        handler = LogtailHandler(source_token=self.source_token, host=self.host, timeout=10)
+        self.assertEqual(handler.uploader.timeout, 10)
 
     @patch('logtail.handler.FlushWorker')
     def test_handler_creates_pipe_from_args(self, MockWorker):

tests/test_uploader.py

@@ -13,10 +13,11 @@ class TestUploader(unittest.TestCase):
     host = 'https://in.logtail.com'
     source_token = 'dummy_source_token'
     frame = [1, 2, 3]
+    timeout = 30
 
     @patch('logtail.uploader.requests.Session.post')
     def test_call(self, post):
-        def mock_post(endpoint, data=None, headers=None):
+        def mock_post(endpoint, data=None, headers=None, timeout=None):
             # Check that the data is sent to ther correct endpoint
             self.assertEqual(endpoint, self.host)
             # Check the content-type
@@ -25,9 +26,11 @@ def mock_post(endpoint, data=None, headers=None):
             self.assertEqual('application/msgpack', headers.get('Content-Type'))
             # Check the content was msgpacked correctly
             self.assertEqual(msgpack.unpackb(data, raw=False), self.frame)
+            # Check that timeout is passed to the request
+            self.assertEqual(timeout, 30)
 
         post.side_effect = mock_post
-        u = Uploader(self.source_token, self.host)
+        u = Uploader(self.source_token, self.host, self.timeout)
         u(self.frame)
 
         self.assertTrue(post.called)