Merge pull request #2 from wealthsimple/master

Fix formatter not being used.

.gitignore

@@ -0,0 +1,2 @@
+*.pyc
+

logzio/handler.py

@@ -4,6 +4,7 @@ import requests
 import traceback
 import datetime
 import json
+import os
 
 from threading import Event, Thread, Condition, Lock, enumerate
 from time import sleep
@@ -101,15 +102,17 @@ class LogzioHandler(logging.Handler):
             # Release the condition
             self.logs_counter_condition.release()
 
-    def handle_exceptions(self, message):
-        if message.exc_info:
-            return '\n'.join(traceback.format_exception(*message.exc_info))
-        else:
-            return message.getMessage()
+    def format(self, record):
+        message = super(LogzioHandler, self).format(record)
+        try:
+            return json.loads(message)
+        except (TypeError, ValueError):
+            return message
 
-    def format_message(self, message):
+    def formatException(self, exc_info):
+        return '\n'.join(traceback.format_exception(*exc_info))
 
-        message_field = self.handle_exceptions(message)
+    def format_message(self, message):
         now = datetime.datetime.utcnow()
         timestamp = now.strftime("%Y-%m-%dT%H:%M:%S") + ".%03d" % (now.microsecond / 1000) + "Z"
 
@@ -118,11 +121,19 @@ class LogzioHandler(logging.Handler):
             "line_number": message.lineno,
             "path_name": message.pathname,
             "log_level": message.levelname,
-            "message": message_field,
             "type": self.logzio_type,
             "@timestamp": timestamp
         }
 
+        if message.exc_info:
+            return_json["message"] = self.formatException(message.exc_info)
+        else:
+            formatted_message = self.format(message)
+            if isinstance(formatted_message, dict):
+                return_json.update(formatted_message)
+            else:
+                return_json["message"] = formatted_message
+
         return return_json
 
     def backup_logs(self, logs):

setup.py

@@ -11,7 +11,7 @@ setup(
     author_email="roi@logz.io",
     url="https://github.com/logzio/logzio-python-handler/",
     license="Apache License 2",
-    packages=find_packages(),
+    packages=find_packages(exclude=["tests"]),
     install_requires=[
         "requests"
     ],

tests/__init__.py

@@ -0,0 +1,12 @@
+import os
+import pkgutil
+
+def load_tests(loader, suite, pattern):
+    for imp, modname, _ in pkgutil.walk_packages(__path__):
+        mod = imp.find_module(modname).load_module(modname)
+        for test in loader.loadTestsFromModule(mod):
+            print("Running TestCase: {}".format(modname))
+            suite.addTests(test)
+
+    print("=" * 70)
+    return suite

tests/unit/handler_test.py

@@ -0,0 +1,119 @@
+import ast
+import unittest
+import logging
+import sys
+import re
+import os
+from logzio.handler import LogzioHandler
+
+def dummy_drain_messages():
+    pass
+
+class TestHandler(unittest.TestCase):
+    def setUp(self):
+        self.handler = LogzioHandler('moo')
+        self.handler.drain_messages = dummy_drain_messages;
+
+    def test_json(self):
+        formatter = logging.Formatter(
+            '{ "appname":"%(name)s", "functionName":"%(funcName)s", \"lineNo":"%(lineno)d", "severity":"%(levelname)s", "message":"%(message)s"}')
+        self.handler.setFormatter(formatter)
+
+        record = logging.LogRecord(
+            name='my-logger',
+            level=0,
+            pathname='handler_test.py',
+            lineno=10,
+            msg="this is a test: moo.",
+            args=[],
+            exc_info=None,
+            func='test_json'
+        )
+
+        formatted_message = self.handler.format_message(record)
+        formatted_message["@timestamp"] = None
+
+        self.assertDictEqual(
+            formatted_message,
+            {
+                '@timestamp': None,
+                'appname': 'my-logger',
+                'functionName': 'test_json',
+                'lineNo': '10',
+                'line_number': 10,
+                'log_level': 'NOTSET',
+                'logger': 'my-logger',
+                'message': 'this is a test: moo.',
+                'path_name': 'handler_test.py',
+                'severity': 'NOTSET',
+                'type': 'python'
+            }
+        )
+
+    def test_string(self):
+        record = logging.LogRecord(
+            name='my-logger',
+            level=0,
+            pathname='handler_test.py',
+            lineno=10,
+            msg="this is a test: moo.",
+            args=[],
+            exc_info=None,
+            func='test_json'
+        )
+
+        formatted_message = self.handler.format_message(record)
+        formatted_message["@timestamp"] = None
+
+        self.assertDictEqual(
+            formatted_message,
+            {
+                '@timestamp': None,
+                'line_number': 10,
+                'log_level': 'NOTSET',
+                'logger': 'my-logger',
+                'message': 'this is a test: moo.',
+                'path_name': 'handler_test.py',
+                'type': 'python'
+             }
+        )
+
+    def test_exc(self):
+        exc_info = None
+        try:
+            raise ValueError("oops.")
+        except:
+            exc_info = sys.exc_info()
+
+        record = logging.LogRecord(
+            name='my-logger',
+            level=0,
+            pathname='handler_test.py',
+            lineno=10,
+            msg="this is a test: moo.",
+            args=[],
+            exc_info=exc_info,
+            func='test_json'
+        )
+
+        formatted_message = self.handler.format_message(record)
+        formatted_message["@timestamp"] = None
+
+        formatted_message["message"] = formatted_message["message"].replace(os.path.abspath(__file__), "")
+        formatted_message["message"] = re.sub(r", line \d+", "", formatted_message["message"])
+
+        self.assertDictEqual(
+            {
+                '@timestamp': None,
+                'line_number': 10,
+                'log_level': 'NOTSET',
+                'logger': 'my-logger',
+                'message': 'Traceback (most recent call last):\n\n  File "", in test_exc\n    raise ValueError("oops.")\n\nValueError: oops.\n',
+                'path_name': 'handler_test.py',
+                'type': 'python'
+            },
+            formatted_message
+        )
+
+if __name__ == '__main__':
+    unittest.main()