T-9413 Add configurable timeout

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)

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

@@ -16,7 +16,7 @@ class TestUploader(unittest.TestCase):
 
     @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,6 +25,8 @@ 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)