Merge pull request #6 from FireTail-io/execution-time

Execution time

Author: rileyfiretail

firetail_lambda/__init__.py

@@ -9,15 +9,16 @@ class firetail_app:
     def __init__(self):
         self.sanitization_callback=default_sanitization_callback
         self.enable_sleeper=False
-        self.sanitize_headers = ['authorization', 'x-api-key']
+        self.sanitize_headers=['authorization', 'x-api-key']
+        self.sleep_time=500
 
 def default_sanitization_callback(event, response):
     return event, response
 
 def firetail_handler(self):
     def decorator(func):
         def wrapper_func(*args, **kwargs):
-            start_time = time.time()
+            wrapper_start_time = time.time()
 
             # make sure it is a valid handler function
             if len(args) < 2:
@@ -27,16 +28,22 @@ def wrapper_func(*args, **kwargs):
             event, _ = args
 
             # Get the response returned down the chain
+            handler_start_time = time.time()
             response = func(*args, **kwargs)
+            handler_execution_time = (time.time() - handler_start_time) * 1000 # Milliseconds
 
-            # Create our log payload, and print it
+            # Sanitise the event & response, create our log payload, and print it
             event, response = self.sanitization_callback(event, response)
-            log_payload = base64.b64encode(json.dumps({"event": event,"response": response}).encode("utf-8")).decode("ascii")
+            log_payload = base64.b64encode(json.dumps({
+                "event": event,
+                "response": response,
+                "execution_time": handler_execution_time
+            }).encode("utf-8")).decode("ascii")
             print("firetail:log-ext:%s" % (log_payload))
 
-            ## Ensure the execution time is >25ms to give the logs API time to propagate our print() to the extension.
+            ## Ensure the execution time is >500ms to give the logs API time to propagate our print() to the extension.
             if self.enable_sleeper:
-                time.sleep(max(time.time() - start_time + 500/1000, 0))
+                time.sleep(max(self.sleep_time/1000 - (time.time() - wrapper_start_time), 0))
 
             # Return the response from down the chain
             return response

tests/test_firetail.py

@@ -1,5 +1,4 @@
-
-
+import base64
 import contextlib
 import copy
 import json
@@ -137,6 +136,15 @@
     }
     }
 
+def rewrite_execution_time(self, output):
+    output = output.split("firetail:log-ext:")[1]
+    response = json.loads(base64.b64decode(output).decode('utf-8'))
+    self.assertIn('execution_time', response)
+    response['execution_time'] = 0.02
+    log_payload = base64.b64encode(json.dumps(response).encode("utf-8")).decode("ascii")
+    response = ("firetail:log-ext:%s" % (log_payload))
+    return response
+
 class TestSimple(unittest.TestCase):
 
     def test_handler_api(self):
@@ -150,8 +158,9 @@ def handler(event, context):
         with contextlib.redirect_stdout(temp_stdout):
             handler(event, "")
         output = temp_stdout.getvalue().strip()
-        self.assertEqual(output, 'firetail:log-ext:eyJldmVudCI6IHt9LCAicmVzcG9uc2UiOiBbMjAxLCAie1wibWVzc2FnZVwiOiBcInN1Y2Nlc3NcIn0iXX0=')
-
+        output = rewrite_execution_time(self, output)
+        self.assertEqual(output, 'firetail:log-ext:eyJldmVudCI6IHt9LCAicmVzcG9uc2UiOiBbMjAxLCAie1wibWVzc2FnZVwiOiBcInN1Y2Nlc3NcIn0iXSwgImV4ZWN1dGlvbl90aW1lIjogMC4wMn0=')
+        
     def test_incorrect_handler_api(self):
         event = {}
         app = firetail_app()
@@ -164,7 +173,6 @@ def handler(argument):
             handler(event)
         output = temp_stdout.getvalue().strip()
         self.assertEqual(output, '')
-
     def test_handler_sleeper_api(self):
         event = {}
         app = firetail_app()
@@ -182,7 +190,8 @@ def handler(event, context):
             difference = end - start
         output = temp_stdout.getvalue().strip()
         self.assertGreaterEqual(difference, .5)
-        self.assertEqual(output, 'firetail:log-ext:eyJldmVudCI6IHt9LCAicmVzcG9uc2UiOiBbMjAxLCAie1wibWVzc2FnZVwiOiBcInN1Y2Nlc3NcIn0iXX0=')
+        output = rewrite_execution_time(self, output)
+        self.assertEqual(output, "firetail:log-ext:eyJldmVudCI6IHt9LCAicmVzcG9uc2UiOiBbMjAxLCAie1wibWVzc2FnZVwiOiBcInN1Y2Nlc3NcIn0iXSwgImV4ZWN1dGlvbl90aW1lIjogMC4wMn0=")
 
     def test_handler_sanitizer(self):
         event = api_gateway_v1
@@ -208,7 +217,8 @@ def handler(event, context):
             handler(event, "")
 
         output = temp_stdout.getvalue().strip()
-        self.assertEqual(output, 'firetail:log-ext:eyJldmVudCI6IHsiYm9keSI6ICJleUowWlhOMElqb2lZbTlrZVNKOSIsICJyZXNvdXJjZSI6ICIve3Byb3h5K30iLCAicGF0aCI6ICIvcGF0aC90by9yZXNvdXJjZSIsICJodHRwTWV0aG9kIjogIlBPU1QiLCAiaXNCYXNlNjRFbmNvZGVkIjogdHJ1ZSwgInF1ZXJ5U3RyaW5nUGFyYW1ldGVycyI6IHsiZm9vIjogImJhciJ9LCAibXVsdGlWYWx1ZVF1ZXJ5U3RyaW5nUGFyYW1ldGVycyI6IHsiZm9vIjogWyJiYXIiXX0sICJwYXRoUGFyYW1ldGVycyI6IHsicHJveHkiOiAiL3BhdGgvdG8vcmVzb3VyY2UifSwgInN0YWdlVmFyaWFibGVzIjogeyJiYXoiOiAicXV4In0sICJoZWFkZXJzIjogeyJBY2NlcHQiOiAidGV4dC9odG1sLGFwcGxpY2F0aW9uL3hodG1sK3htbCxhcHBsaWNhdGlvbi94bWw7cT0wLjksaW1hZ2Uvd2VicCwqLyo7cT0wLjgiLCAiQWNjZXB0LUVuY29kaW5nIjogImd6aXAsIGRlZmxhdGUsIHNkY2giLCAiQWNjZXB0LUxhbmd1YWdlIjogImVuLVVTLGVuO3E9MC44IiwgIkNhY2hlLUNvbnRyb2wiOiAibWF4LWFnZT0wIiwgIkNsb3VkRnJvbnQtRm9yd2FyZGVkLVByb3RvIjogImh0dHBzIiwgIkNsb3VkRnJvbnQtSXMtRGVza3RvcC1WaWV3ZXIiOiAidHJ1ZSIsICJDbG91ZEZyb250LUlzLU1vYmlsZS1WaWV3ZXIiOiAiZmFsc2UiLCAiQ2xvdWRGcm9udC1Jcy1TbWFydFRWLVZpZXdlciI6ICJmYWxzZSIsICJDbG91ZEZyb250LUlzLVRhYmxldC1WaWV3ZXIiOiAiZmFsc2UiLCAiQ2xvdWRGcm9udC1WaWV3ZXItQ291bnRyeSI6ICJVUyIsICJIb3N0IjogIjEyMzQ1Njc4OTAuZXhlY3V0ZS1hcGkudXMtZWFzdC0xLmFtYXpvbmF3cy5jb20iLCAiVXBncmFkZS1JbnNlY3VyZS1SZXF1ZXN0cyI6ICIxIiwgIlVzZXItQWdlbnQiOiAiQ3VzdG9tIFVzZXIgQWdlbnQgU3RyaW5nIiwgIlZpYSI6ICIxLjEgMDhmMzIzZGVhZGJlZWZhN2FmMzRkNWZlYjQxNGNlMjcuY2xvdWRmcm9udC5uZXQgKENsb3VkRnJvbnQpIiwgIlgtQW16LUNmLUlkIjogImNEZWhWUW9abng0M1ZZUWI5ajItbnZDaC05ejM5NlVoYnAwMjdZMkp2a0NQTkxtR0pIcWxhQT09IiwgIlgtRm9yd2FyZGVkLUZvciI6ICIxMjcuMC4wLjEsIDEyNy4wLjAuMiIsICJYLUZvcndhcmRlZC1Qb3J0IjogIjQ0MyIsICJYLUZvcndhcmRlZC1Qcm90byI6ICJodHRwcyJ9LCAibXVsdGlWYWx1ZUhlYWRlcnMiOiB7IkFjY2VwdCI6IFsidGV4dC9odG1sLGFwcGxpY2F0aW9uL3hodG1sK3htbCxhcHBsaWNhdGlvbi94bWw7cT0wLjksaW1hZ2Uvd2VicCwqLyo7cT0wLjgiXSwgIkFjY2VwdC1FbmNvZGluZyI6IFsiZ3ppcCwgZGVmbGF0ZSwgc2RjaCJdLCAiQWNjZXB0LUxhbmd1YWdlIjogWyJlbi1VUyxlbjtxPTAuOCJdLCAiQ2FjaGUtQ29udHJvbCI6IFsibWF4LWFnZT0wIl0sICJDbG91ZEZyb250LUZvcndhcmRlZC1Qcm90byI6IFsiaHR0cHMiXSwgIkNsb3VkRnJvbnQtSXMtRGVza3RvcC1WaWV3ZXIiOiBbInRydWUiXSwgIkNsb3VkRnJvbnQtSXMtTW9iaWxlLVZpZXdlciI6IFsiZmFsc2UiXSwgIkNsb3VkRnJvbnQtSXMtU21hcnRUVi1WaWV3ZXIiOiBbImZhbHNlIl0sICJDbG91ZEZyb250LUlzLVRhYmxldC1WaWV3ZXIiOiBbImZhbHNlIl0sICJDbG91ZEZyb250LVZpZXdlci1Db3VudHJ5IjogWyJVUyJdLCAiSG9zdCI6IFsiMDEyMzQ1Njc4OS5leGVjdXRlLWFwaS51cy1lYXN0LTEuYW1hem9uYXdzLmNvbSJdLCAiVXBncmFkZS1JbnNlY3VyZS1SZXF1ZXN0cyI6IFsiMSJdLCAiVXNlci1BZ2VudCI6IFsiQ3VzdG9tIFVzZXIgQWdlbnQgU3RyaW5nIl0sICJWaWEiOiBbIjEuMSAwOGYzMjNkZWFkYmVlZmE3YWYzNGQ1ZmViNDE0Y2UyNy5jbG91ZGZyb250Lm5ldCAoQ2xvdWRGcm9udCkiXSwgIlgtQW16LUNmLUlkIjogWyJjRGVoVlFvWm54NDNWWVFiOWoyLW52Q2gtOXozOTZVaGJwMDI3WTJKdmtDUE5MbUdKSHFsYUE9PSJdLCAiWC1Gb3J3YXJkZWQtRm9yIjogWyIxMjcuMC4wLjEsIDEyNy4wLjAuMiJdLCAiWC1Gb3J3YXJkZWQtUG9ydCI6IFsiNDQzIl0sICJYLUZvcndhcmRlZC1Qcm90byI6IFsiaHR0cHMiXX0sICJyZXF1ZXN0Q29udGV4dCI6IHsiYWNjb3VudElkIjogIjEyMzQ1Njc4OTAxMiIsICJyZXNvdXJjZUlkIjogIjEyMzQ1NiIsICJzdGFnZSI6ICJwcm9kIiwgInJlcXVlc3RJZCI6ICJjNmFmOWFjNi03YjYxLTExZTYtOWE0MS05M2U4ZGVhZGJlZWYiLCAicmVxdWVzdFRpbWUiOiAiMDkvQXByLzIwMTU6MTI6MzQ6NTYgKzAwMDAiLCAicmVxdWVzdFRpbWVFcG9jaCI6IDE0Mjg1ODI4OTYwMDAsICJpZGVudGl0eSI6IHsiY29nbml0b0lkZW50aXR5UG9vbElkIjogbnVsbCwgImFjY291bnRJZCI6IG51bGwsICJjb2duaXRvSWRlbnRpdHlJZCI6IG51bGwsICJjYWxsZXIiOiBudWxsLCAiYWNjZXNzS2V5IjogbnVsbCwgInNvdXJjZUlwIjogIjEyNy4wLjAuMSIsICJjb2duaXRvQXV0aGVudGljYXRpb25UeXBlIjogbnVsbCwgImNvZ25pdG9BdXRoZW50aWNhdGlvblByb3ZpZGVyIjogbnVsbCwgInVzZXJBcm4iOiBudWxsLCAidXNlckFnZW50IjogIkN1c3RvbSBVc2VyIEFnZW50IFN0cmluZyIsICJ1c2VyIjogbnVsbH0sICJwYXRoIjogIi9wcm9kL3BhdGgvdG8vcmVzb3VyY2UiLCAicmVzb3VyY2VQYXRoIjogIi97cHJveHkrfSIsICJodHRwTWV0aG9kIjogIlBPU1QiLCAiYXBpSWQiOiAiMTIzNDU2Nzg5MCIsICJwcm90b2NvbCI6ICJIVFRQLzEuMSJ9fSwgInJlc3BvbnNlIjogWzIwMSwgIntcIm1lc3NhZ2VcIjogXCJzdWNjZXNzXCJ9Il19')
+        output = rewrite_execution_time(self, output)
+        self.assertEqual(output, 'firetail:log-ext:eyJldmVudCI6IHsiYm9keSI6ICJleUowWlhOMElqb2lZbTlrZVNKOSIsICJyZXNvdXJjZSI6ICIve3Byb3h5K30iLCAicGF0aCI6ICIvcGF0aC90by9yZXNvdXJjZSIsICJodHRwTWV0aG9kIjogIlBPU1QiLCAiaXNCYXNlNjRFbmNvZGVkIjogdHJ1ZSwgInF1ZXJ5U3RyaW5nUGFyYW1ldGVycyI6IHsiZm9vIjogImJhciJ9LCAibXVsdGlWYWx1ZVF1ZXJ5U3RyaW5nUGFyYW1ldGVycyI6IHsiZm9vIjogWyJiYXIiXX0sICJwYXRoUGFyYW1ldGVycyI6IHsicHJveHkiOiAiL3BhdGgvdG8vcmVzb3VyY2UifSwgInN0YWdlVmFyaWFibGVzIjogeyJiYXoiOiAicXV4In0sICJoZWFkZXJzIjogeyJBY2NlcHQiOiAidGV4dC9odG1sLGFwcGxpY2F0aW9uL3hodG1sK3htbCxhcHBsaWNhdGlvbi94bWw7cT0wLjksaW1hZ2Uvd2VicCwqLyo7cT0wLjgiLCAiQWNjZXB0LUVuY29kaW5nIjogImd6aXAsIGRlZmxhdGUsIHNkY2giLCAiQWNjZXB0LUxhbmd1YWdlIjogImVuLVVTLGVuO3E9MC44IiwgIkNhY2hlLUNvbnRyb2wiOiAibWF4LWFnZT0wIiwgIkNsb3VkRnJvbnQtRm9yd2FyZGVkLVByb3RvIjogImh0dHBzIiwgIkNsb3VkRnJvbnQtSXMtRGVza3RvcC1WaWV3ZXIiOiAidHJ1ZSIsICJDbG91ZEZyb250LUlzLU1vYmlsZS1WaWV3ZXIiOiAiZmFsc2UiLCAiQ2xvdWRGcm9udC1Jcy1TbWFydFRWLVZpZXdlciI6ICJmYWxzZSIsICJDbG91ZEZyb250LUlzLVRhYmxldC1WaWV3ZXIiOiAiZmFsc2UiLCAiQ2xvdWRGcm9udC1WaWV3ZXItQ291bnRyeSI6ICJVUyIsICJIb3N0IjogIjEyMzQ1Njc4OTAuZXhlY3V0ZS1hcGkudXMtZWFzdC0xLmFtYXpvbmF3cy5jb20iLCAiVXBncmFkZS1JbnNlY3VyZS1SZXF1ZXN0cyI6ICIxIiwgIlVzZXItQWdlbnQiOiAiQ3VzdG9tIFVzZXIgQWdlbnQgU3RyaW5nIiwgIlZpYSI6ICIxLjEgMDhmMzIzZGVhZGJlZWZhN2FmMzRkNWZlYjQxNGNlMjcuY2xvdWRmcm9udC5uZXQgKENsb3VkRnJvbnQpIiwgIlgtQW16LUNmLUlkIjogImNEZWhWUW9abng0M1ZZUWI5ajItbnZDaC05ejM5NlVoYnAwMjdZMkp2a0NQTkxtR0pIcWxhQT09IiwgIlgtRm9yd2FyZGVkLUZvciI6ICIxMjcuMC4wLjEsIDEyNy4wLjAuMiIsICJYLUZvcndhcmRlZC1Qb3J0IjogIjQ0MyIsICJYLUZvcndhcmRlZC1Qcm90byI6ICJodHRwcyJ9LCAibXVsdGlWYWx1ZUhlYWRlcnMiOiB7IkFjY2VwdCI6IFsidGV4dC9odG1sLGFwcGxpY2F0aW9uL3hodG1sK3htbCxhcHBsaWNhdGlvbi94bWw7cT0wLjksaW1hZ2Uvd2VicCwqLyo7cT0wLjgiXSwgIkFjY2VwdC1FbmNvZGluZyI6IFsiZ3ppcCwgZGVmbGF0ZSwgc2RjaCJdLCAiQWNjZXB0LUxhbmd1YWdlIjogWyJlbi1VUyxlbjtxPTAuOCJdLCAiQ2FjaGUtQ29udHJvbCI6IFsibWF4LWFnZT0wIl0sICJDbG91ZEZyb250LUZvcndhcmRlZC1Qcm90byI6IFsiaHR0cHMiXSwgIkNsb3VkRnJvbnQtSXMtRGVza3RvcC1WaWV3ZXIiOiBbInRydWUiXSwgIkNsb3VkRnJvbnQtSXMtTW9iaWxlLVZpZXdlciI6IFsiZmFsc2UiXSwgIkNsb3VkRnJvbnQtSXMtU21hcnRUVi1WaWV3ZXIiOiBbImZhbHNlIl0sICJDbG91ZEZyb250LUlzLVRhYmxldC1WaWV3ZXIiOiBbImZhbHNlIl0sICJDbG91ZEZyb250LVZpZXdlci1Db3VudHJ5IjogWyJVUyJdLCAiSG9zdCI6IFsiMDEyMzQ1Njc4OS5leGVjdXRlLWFwaS51cy1lYXN0LTEuYW1hem9uYXdzLmNvbSJdLCAiVXBncmFkZS1JbnNlY3VyZS1SZXF1ZXN0cyI6IFsiMSJdLCAiVXNlci1BZ2VudCI6IFsiQ3VzdG9tIFVzZXIgQWdlbnQgU3RyaW5nIl0sICJWaWEiOiBbIjEuMSAwOGYzMjNkZWFkYmVlZmE3YWYzNGQ1ZmViNDE0Y2UyNy5jbG91ZGZyb250Lm5ldCAoQ2xvdWRGcm9udCkiXSwgIlgtQW16LUNmLUlkIjogWyJjRGVoVlFvWm54NDNWWVFiOWoyLW52Q2gtOXozOTZVaGJwMDI3WTJKdmtDUE5MbUdKSHFsYUE9PSJdLCAiWC1Gb3J3YXJkZWQtRm9yIjogWyIxMjcuMC4wLjEsIDEyNy4wLjAuMiJdLCAiWC1Gb3J3YXJkZWQtUG9ydCI6IFsiNDQzIl0sICJYLUZvcndhcmRlZC1Qcm90byI6IFsiaHR0cHMiXX0sICJyZXF1ZXN0Q29udGV4dCI6IHsiYWNjb3VudElkIjogIjEyMzQ1Njc4OTAxMiIsICJyZXNvdXJjZUlkIjogIjEyMzQ1NiIsICJzdGFnZSI6ICJwcm9kIiwgInJlcXVlc3RJZCI6ICJjNmFmOWFjNi03YjYxLTExZTYtOWE0MS05M2U4ZGVhZGJlZWYiLCAicmVxdWVzdFRpbWUiOiAiMDkvQXByLzIwMTU6MTI6MzQ6NTYgKzAwMDAiLCAicmVxdWVzdFRpbWVFcG9jaCI6IDE0Mjg1ODI4OTYwMDAsICJpZGVudGl0eSI6IHsiY29nbml0b0lkZW50aXR5UG9vbElkIjogbnVsbCwgImFjY291bnRJZCI6IG51bGwsICJjb2duaXRvSWRlbnRpdHlJZCI6IG51bGwsICJjYWxsZXIiOiBudWxsLCAiYWNjZXNzS2V5IjogbnVsbCwgInNvdXJjZUlwIjogIjEyNy4wLjAuMSIsICJjb2duaXRvQXV0aGVudGljYXRpb25UeXBlIjogbnVsbCwgImNvZ25pdG9BdXRoZW50aWNhdGlvblByb3ZpZGVyIjogbnVsbCwgInVzZXJBcm4iOiBudWxsLCAidXNlckFnZW50IjogIkN1c3RvbSBVc2VyIEFnZW50IFN0cmluZyIsICJ1c2VyIjogbnVsbH0sICJwYXRoIjogIi9wcm9kL3BhdGgvdG8vcmVzb3VyY2UiLCAicmVzb3VyY2VQYXRoIjogIi97cHJveHkrfSIsICJodHRwTWV0aG9kIjogIlBPU1QiLCAiYXBpSWQiOiAiMTIzNDU2Nzg5MCIsICJwcm90b2NvbCI6ICJIVFRQLzEuMSJ9fSwgInJlc3BvbnNlIjogWzIwMSwgIntcIm1lc3NhZ2VcIjogXCJzdWNjZXNzXCJ9Il0sICJleGVjdXRpb25fdGltZSI6IDAuMDJ9')
 
 
 if __name__ == '__main__': # pragma: no cover