Merge pull request #2973 from bblommers/master

Lambda - Return LogResult only if requested

moto/awslambda/models.py

@@ -381,6 +381,7 @@ class LambdaFunction(BaseModel):
             event = dict()
         if context is None:
             context = {}
+        output = None
 
         try:
             # TODO: I believe we can keep the container running and feed events as needed
@@ -396,7 +397,7 @@ class LambdaFunction(BaseModel):
 
             env_vars.update(self.environment_vars)
 
-            container = output = exit_code = None
+            container = exit_code = None
             log_config = docker.types.LogConfig(type=docker.types.LogConfig.types.JSON)
             with _DockerDataVolumeContext(self) as data_vol:
                 try:
@@ -457,24 +458,28 @@ class LambdaFunction(BaseModel):
 
             # We only care about the response from the lambda
             # Which is the last line of the output, according to https://github.com/lambci/docker-lambda/issues/25
-            output = output.splitlines()[-1]
+            resp = output.splitlines()[-1]
-            return output, False
+            logs = os.linesep.join(
+                [line for line in self.convert(output).splitlines()[:-1]]
+            )
+            return resp, False, logs
         except BaseException as e:
             traceback.print_exc()
-            return "error running lambda: {}".format(e), True
+            logs = os.linesep.join(
+                [line for line in self.convert(output).splitlines()[:-1]]
+            )
+            return "error running lambda: {}".format(e), True, logs
 
     def invoke(self, body, request_headers, response_headers):
-        payload = dict()
 
         if body:
             body = json.loads(body)
 
         # Get the invocation type:
-        res, errored = self._invoke_lambda(code=self.code, event=body)
+        res, errored, logs = self._invoke_lambda(code=self.code, event=body)
         if request_headers.get("x-amz-invocation-type") == "RequestResponse":
-            encoded = base64.b64encode(res.encode("utf-8"))
+            encoded = base64.b64encode(logs.encode("utf-8"))
             response_headers["x-amz-log-result"] = encoded.decode("utf-8")
-            payload["result"] = response_headers["x-amz-log-result"]
             result = res.encode("utf-8")
         else:
             result = res
@@ -1068,9 +1073,9 @@ class LambdaBackend(BaseBackend):
         if fn:
             payload = fn.invoke(body, headers, response_headers)
             response_headers["Content-Length"] = str(len(payload))
-            return response_headers, payload
+            return payload
         else:
-            return response_headers, None
+            return None
 
 
 def do_validate_s3():

moto/awslambda/responses.py

@@ -178,7 +178,7 @@ class LambdaResponse(BaseResponse):
         function_name = unquote(self.path.rsplit("/", 2)[-2])
         qualifier = self._get_param("qualifier")
 
-        response_header, payload = self.lambda_backend.invoke(
+        payload = self.lambda_backend.invoke(
             function_name, qualifier, self.body, self.headers, response_headers
         )
         if payload:
@@ -187,6 +187,8 @@ class LambdaResponse(BaseResponse):
             elif request.headers.get("X-Amz-Invocation-Type") == "DryRun":
                 status_code = 204
             else:
+                if request.headers.get("X-Amz-Log-Type") != "Tail":
+                    del response_headers["x-amz-log-result"]
                 status_code = 200
             return status_code, response_headers, payload
         else:

tests/test_awslambda/test_lambda.py

@@ -43,6 +43,7 @@ def _process_lambda(func_str):
 def get_test_zip_file1():
     pfunc = """
 def lambda_handler(event, context):
+    print("custom log event")
     return event
 """
     return _process_lambda(pfunc)
@@ -112,18 +113,29 @@ def test_invoke_requestresponse_function():
         FunctionName="testFunction",
         InvocationType="RequestResponse",
         Payload=json.dumps(in_data),
+        LogType="Tail",
     )
 
     success_result["StatusCode"].should.equal(200)
-    result_obj = json.loads(
+    logs = base64.b64decode(success_result["LogResult"]).decode("utf-8")
-        base64.b64decode(success_result["LogResult"]).decode("utf-8")
-    )
 
-    result_obj.should.equal(in_data)
+    logs.should.contain("START RequestId:")
+    logs.should.contain("custom log event")
+    logs.should.contain("END RequestId:")
 
     payload = success_result["Payload"].read().decode("utf-8")
     json.loads(payload).should.equal(in_data)
 
+    # Logs should not be returned by default, only when the LogType-param is supplied
+    success_result = conn.invoke(
+        FunctionName="testFunction",
+        InvocationType="RequestResponse",
+        Payload=json.dumps(in_data),
+    )
+
+    success_result["StatusCode"].should.equal(200)
+    assert "LogResult" not in success_result
+
 
 @mock_lambda
 def test_invoke_requestresponse_function_with_arn():
@@ -152,11 +164,6 @@ def test_invoke_requestresponse_function_with_arn():
     )
 
     success_result["StatusCode"].should.equal(200)
-    result_obj = json.loads(
-        base64.b64decode(success_result["LogResult"]).decode("utf-8")
-    )
-
-    result_obj.should.equal(in_data)
 
     payload = success_result["Payload"].read().decode("utf-8")
     json.loads(payload).should.equal(in_data)