Revise logic and clean up tests

moesif_aws_lambda/middleware.py

@@ -198,39 +198,27 @@ def base64_body(cls, data):
             else:
                 return str(body), 'base64'
 
-        def safe_json_parse(self, body, is_response):
+        def safe_json_parse(self, body):
             """Tries to parse the `body` as JSON safely.
             Returns the formatted body and the appropriate `transfer_encoding`. 
             """
             try:
                 if isinstance(body, (dict, list)):
                     # If body is an instance of either a dictionary of list, 
                     # we can return it as is.
-                    return body, None
-                """
-                If body is neither dictionary or list, it has to be one of these types: 
-                - binary data (`bytes` object in Python)
-                - a string
-                - a non-string type object like an integer or float
-                """
-                body_str = None
-                if isinstance(body, bytes):
+                    return body, "json"
+                elif isinstance(body, bytes):
                     body_str = body.decode()
+                    parsed_body = json.loads(body_str)
+                    return parsed_body, "json"
                 else:
-                    body_str = str(body)
-
-                # Now we try to parse the string as JSON
-                json_body = json.loads(body_str)
-                if not is_response:
-                    return json_body, None 
-                else:
-                    return json_body, "json"
-
-
+                    parsed_body = json.loads(body)
+                    return parsed_body, "json"
+
             except (json.JSONDecodeError, TypeError, ValueError, UnicodeError) as error:
                 return self.base64_body(body)
 
-        def process_body(self, body_wrapper, is_response):
+        def process_body(self, body_wrapper):
             """Function to process body"""
 
             if self.LOG_BODY and isinstance(body_wrapper, dict) and 'body' not in body_wrapper:
@@ -246,11 +234,11 @@ def process_body(self, body_wrapper, is_response):
             transfer_encoding = None
 
             try:
-                if body_wrapper.get('isBase64Encoded', False) and (is_response or self.is_base64_str(body_wrapper.get('body'))):
-                    body = body_wrapper.get('body')
-                    transfer_encoding = 'base64'
+                if body_wrapper.get('isBase64Encoded', False) and self.is_base64_str(body_wrapper.get('body')):
+                        body = body_wrapper.get('body')
+                        transfer_encoding = 'base64'
                 else:
-                    body, transfer_encoding = self.safe_json_parse(body_wrapper.get('body'), is_response)
+                    body, transfer_encoding = self.safe_json_parse(body_wrapper.get('body'))
             except Exception as e:
                     return self.base64_body(body_wrapper['body'])
 
@@ -309,7 +297,7 @@ def before(self, event, context):
                 request_time = datetime.utcnow()
 
             # Request Body
-            req_body, req_transfer_encoding = self.process_body(event, False)
+            req_body, req_transfer_encoding = self.process_body(event)
 
             # Metadata
             start_time_get_metadata = datetime.utcnow()
@@ -417,7 +405,7 @@ def after(self, retval):
             event_send = None
             if self.event is not None:
                 # Response body
-                resp_body, resp_transfer_encoding = self.process_body(retval, True)
+                resp_body, resp_transfer_encoding = self.process_body(retval)
 
                 # Event Response object
                 event_rsp = EventResponseModel(time = datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3],

moesif_aws_lambda/tests/event_body_json.json

@@ -55,7 +55,7 @@
     "routeKey": "$default",
     "stage": "$default"
   },
-  "body": "{'foo': 'bar'}",
+  "body": "{\"foo\": \"bar\"}",
   "pathParameters": {
     "parameter1": "value1"
   },

moesif_aws_lambda/tests/tests.py

@@ -51,7 +51,7 @@ def test_request_with_valid_base64_body(self):
         Moesif = MoesifLogger(moesif_options)
         moesif_middleware = Moesif(lambda_handler)
         _ = moesif_middleware(event_payload, {})
-        req_body, transfer_encoding = moesif_middleware.process_body(event_payload, False)
+        req_body, transfer_encoding = moesif_middleware.process_body(event_payload)
 
         self.assertEqual(req_body, "eyJ0ZXN0IjoiYm9keSJ9")
         self.assertEqual(transfer_encoding, "base64")
@@ -70,9 +70,9 @@ def test_request_with_invalid_base64_body_types(self):
         ]
 
         expected = [
-            ({"foo": "bar"}, None),
-            (json.loads(str(10)), None),
-            ("eydmb28nOiAnYmFyJ30=", "base64"),
+            ({"foo": "bar"}, "json"),
+            ("MTA=", "base64"),
+            ({"foo": "bar"}, "json"),
         ]
 
         for i, file in enumerate(test_files):
@@ -82,7 +82,7 @@ def test_request_with_invalid_base64_body_types(self):
             moesif_middleware = moesif(lambda_handler)
             _ = moesif_middleware(payload, {})
 
-            req_body, transfer_encoding = moesif_middleware.process_body(payload, False)
+            req_body, transfer_encoding = moesif_middleware.process_body(payload)
             self.assertTupleEqual((req_body, transfer_encoding), expected[i])
 
     def test_response_with_json_body(self):
@@ -94,7 +94,7 @@ def test_response_with_json_body(self):
             response_payload = json.load(event)
         Moesif = MoesifLogger(moesif_options)
         moesif_middleware = Moesif(lambda_handler)
-        res_body, transfer_encoding = moesif_middleware.process_body(response_payload, True)
+        res_body, transfer_encoding = moesif_middleware.process_body(response_payload)
 
         expected_body = {"message": "Hello from Lambda!"}
 
@@ -107,24 +107,19 @@ def test_response_with_b64_body(self):
         Body for this test is binary data encoded in base64.
         """
 
-        with open("moesif_aws_lambda/tests/image.png", 'rb') as img_file:
-            img = img_file.read()
-
-        b64_img = base64.b64encode(img).decode("utf-8"),
-
         response_payload = {
             "statusCode": 200,
             "isBase64Encoded": True,
-            "body": b64_img,
-            "headers": {"Content-Type": "image/png"}, 
+            "body": "eyJmb28iOiAiYmFyIn0=",
+            "headers": {"Content-Type": "application/json"}, 
         }
 
         Moesif = MoesifLogger(moesif_options)
         moesif_middleware = Moesif(lambda_handler)
 
-        res_body, transfer_encoding = moesif_middleware.process_body(response_payload, True)
+        res_body, transfer_encoding = moesif_middleware.process_body(response_payload)
 
-        self.assertEqual(res_body, b64_img)
+        self.assertEqual(res_body, "eyJmb28iOiAiYmFyIn0=")
         self.assertEqual(transfer_encoding, "base64")
 
 if __name__ == "__main__":