Use features instead of body as input

lambda_function.py

@@ -4,15 +4,20 @@ import json
 
 cls = Classifier()
 
+# Lambda handler (proxy integration option unchecked on AWS API Gateway)
 def lambda_handler(event, context):
-    try:
-        # Parse the input data
-        data = json.loads(event.get('body', '{}'))
 
-        response = cls.load_and_test(data)
+    try:
+        features = event.get('features', {})
+        if not features:
+            raise ValueError("'features' key missing")
 
+        response = cls.load_and_test(features)
         return {
             'statusCode': 200,
+            'headers': {
+                'Content-Type': 'application/json'
+            },
             'body': json.dumps({
                 'predictions': response["predictions"],
                 'probabilities': response["probabilities"]

tests/test_lambda.py

@@ -11,7 +11,9 @@ from lambda_function import lambda_handler
 
 @pytest.fixture
 def event():
-    json_event = {"body": "{\"features\": [[6.5, 3.0, 5.8, 2.2], [6.1, 2.8, 4.7, 1.2]]}"}
+    json_event = {
+        "features": [[6.5, 3.0, 5.8, 2.2], [6.1, 2.8, 4.7, 1.2]]
+    }
     return json_event
 
 @pytest.fixture