updated test case to cover process

Author: Riley Priddle

firetail_lambda/__init__.py

@@ -4,26 +4,24 @@
 def firetail_handler(enable_sleeper=False):
     def decorator(func):
         def wrapper_func(*args, **kwargs):
-            compatible = True
             start_time = time.time()
 
             # make sure it is a valid handler function
             if len(args) < 2:
-                compatible = False
-            if compatible:
-                # Unpack the args
-                event, _ = args
+                return func(*args, **kwargs)
+
+            # Unpack the args
+            event, _ = args
 
             # Get the response returned down the chain
             response = func(*args, **kwargs)
 
             # Create our log payload, and print it
-            if compatible:
-                log_payload = base64.b64encode(json.dumps({"event": event,"response": response}).encode("utf-8")).decode("ascii")
-                print("firetail:loggingapi:%s" % (log_payload))
+            log_payload = base64.b64encode(json.dumps({"event": event,"response": response}).encode("utf-8")).decode("ascii")
+            print("firetail:loggingapi:%s" % (log_payload))
 
             ## Ensure the execution time is >25ms to give the logs API time to propagate our print() to the extension.
-            if enable_sleeper and compatible:
+            if enable_sleeper:
                 time.sleep(max(time.time() - start_time + 500/1000, 0))
 
             # Return the response from down the chain

tests/test_firetail.py

@@ -0,0 +1,24 @@
+import unittest
+import json
+import contextlib
+from io import StringIO
+
+from firetail_lambda import firetail_handler
+
+class TestSimple(unittest.TestCase):
+
+    def test_handler_api(self):
+        event = {}
+        @firetail_handler()
+        def handler(event, context):
+            return 201, json.dumps({"message": "success"})
+        
+        temp_stdout = StringIO()
+        with contextlib.redirect_stdout(temp_stdout):
+            handler(event, "")
+        output = temp_stdout.getvalue().strip()
+        self.assertEqual(output, 'firetail:loggingapi:eyJldmVudCI6IHt9LCAicmVzcG9uc2UiOiBbMjAxLCAie1wibWVzc2FnZVwiOiBcInN1Y2Nlc3NcIn0iXX0=')
+
+
+if __name__ == '__main__':
+    unittest.main()