From 2ae09c5335c237105f61c12342ae97aa190efd73 Mon Sep 17 00:00:00 2001
From: Roque Pinel <1685896+repinel@users.noreply.github.com>
Date: Tue, 14 Aug 2018 19:00:03 -0400
Subject: [PATCH] Fix the `StatusCode` returned by lambda invoke

According to the AWS documentation:
```
The HTTP status code will be in the 200 range for successful request.
For the RequestResponse invocation type this status code will be 200.
For the Event invocation type this status code will be 202.
For the DryRun invocation type the status code will be 204.
```
---
 moto/awslambda/responses.py         |  8 ++++++-
 tests/test_awslambda/test_lambda.py | 35 +++++++++++++++++++++++++++--
 2 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/moto/awslambda/responses.py b/moto/awslambda/responses.py
index 46203c10d..3d9b3ee3c 100644
--- a/moto/awslambda/responses.py
+++ b/moto/awslambda/responses.py
@@ -172,7 +172,13 @@ class LambdaResponse(BaseResponse):
             function_name, qualifier, self.body, self.headers, response_headers
         )
         if payload:
-            return 202, response_headers, payload
+            if request.headers['X-Amz-Invocation-Type'] == 'Event':
+                status_code = 202
+            elif request.headers['X-Amz-Invocation-Type'] == 'DryRun':
+                status_code = 204
+            else:
+                status_code = 200
+            return status_code, response_headers, payload
         else:
             return 404, response_headers, "{}"
 
diff --git a/tests/test_awslambda/test_lambda.py b/tests/test_awslambda/test_lambda.py
index 2835729f8..6601537fd 100644
--- a/tests/test_awslambda/test_lambda.py
+++ b/tests/test_awslambda/test_lambda.py
@@ -113,7 +113,7 @@ def test_invoke_requestresponse_function():
         Payload=json.dumps(in_data),
     )
 
-    success_result["StatusCode"].should.equal(202)
+    success_result["StatusCode"].should.equal(200)
     result_obj = json.loads(
         base64.b64decode(success_result["LogResult"]).decode("utf-8")
     )
@@ -151,6 +151,37 @@ def test_invoke_event_function():
     json.loads(success_result["Payload"].read().decode("utf-8")).should.equal({})
 
 
+@mock_lambda
+def test_invoke_dryrun_function():
+    conn = boto3.client('lambda', 'us-west-2')
+    conn.create_function(
+        FunctionName='testFunction',
+        Runtime='python2.7',
+        Role=get_role_name(),
+        Handler='lambda_function.lambda_handler',
+        Code={
+            'ZipFile': get_test_zip_file1(),
+        },
+        Description='test lambda function',
+        Timeout=3,
+        MemorySize=128,
+        Publish=True,
+    )
+
+    conn.invoke.when.called_with(
+        FunctionName='notAFunction',
+        InvocationType='Event',
+        Payload='{}'
+    ).should.throw(botocore.client.ClientError)
+
+    in_data = {'msg': 'So long and thanks for all the fish'}
+    success_result = conn.invoke(
+        FunctionName='testFunction', InvocationType='DryRun', Payload=json.dumps(in_data))
+    success_result["StatusCode"].should.equal(204)
+    json.loads(success_result['Payload'].read().decode(
+        'utf-8')).should.equal({})
+
+
 if settings.TEST_SERVER_MODE:
 
     @mock_ec2
@@ -179,7 +210,7 @@ if settings.TEST_SERVER_MODE:
             InvocationType="RequestResponse",
             Payload=json.dumps(in_data),
         )
-        result["StatusCode"].should.equal(202)
+        result["StatusCode"].should.equal(200)
         actual_payload = json.loads(result["Payload"].read().decode("utf-8"))
         expected_payload = {"id": vol.id, "state": vol.state, "size": vol.size}
         actual_payload.should.equal(expected_payload)