From 62a2d8f75692329bac4b694e41a4fdd197dcde76 Mon Sep 17 00:00:00 2001
From: Terry Cain <terry@terrys-home.co.uk>
Date: Wed, 25 Oct 2017 20:04:00 +0100
Subject: [PATCH] Improved lambda timings

---
 moto/awslambda/models.py            |  7 ++++++-
 moto/awslambda/responses.py         | 24 ++++++++++++++++++++++--
 tests/test_awslambda/test_lambda.py | 13 ++++++++-----
 3 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/moto/awslambda/models.py b/moto/awslambda/models.py
index 935abbcd6..6306acd5c 100644
--- a/moto/awslambda/models.py
+++ b/moto/awslambda/models.py
@@ -298,7 +298,12 @@ class LambdaFunction(BaseModel):
                         volumes=["{}:/var/task".format(data_vol.name)], environment=env_vars, detach=True, **run_kwargs)
                 finally:
                     if container:
-                        exit_code = container.wait()
+                        try:
+                            exit_code = container.wait(timeout=300)
+                        except requests.exceptions.ReadTimeout:
+                            exit_code = -1
+                            container.stop()
+                            container.kill()
                         output = container.logs(stdout=False, stderr=True)
                         output += container.logs(stdout=True, stderr=False)
                         container.remove()
diff --git a/moto/awslambda/responses.py b/moto/awslambda/responses.py
index 4ba837ea2..32f9e2e01 100644
--- a/moto/awslambda/responses.py
+++ b/moto/awslambda/responses.py
@@ -12,6 +12,24 @@ except:
 from moto.core.utils import amz_crc32, amzn_request_id
 from moto.core.responses import BaseResponse
 
+import signal
+import contextlib
+
+
+@contextlib.contextmanager
+def timeout(seconds=5):
+    def handle_error(signum, frame):
+        raise RuntimeError()
+
+    try:
+        signal.signal(signal.SIGALRM, handle_error)
+        signal.alarm(seconds)
+        yield None
+    except RuntimeError:
+        pass
+    finally:
+        signal.alarm(0)
+
 
 class LambdaResponse(BaseResponse):
 
@@ -100,7 +118,8 @@ class LambdaResponse(BaseResponse):
 
         if lambda_backend.has_function(function_name):
             fn = lambda_backend.get_function(function_name)
-            payload = fn.invoke(self.body, self.headers, response_headers)
+            with timeout(seconds=360):
+                payload = fn.invoke(self.body, self.headers, response_headers)
             response_headers['Content-Length'] = str(len(payload))
             return 202, response_headers, payload
         else:
@@ -114,7 +133,8 @@ class LambdaResponse(BaseResponse):
         function_name = path.split('/')[-3]
         if lambda_backend.has_function(function_name):
             fn = lambda_backend.get_function(function_name)
-            fn.invoke(self.body, self.headers, response_headers)
+            with timeout(seconds=360):
+                fn.invoke(self.body, self.headers, response_headers)
             response_headers['Content-Length'] = str(0)
             return 202, response_headers, ""
         else:
diff --git a/tests/test_awslambda/test_lambda.py b/tests/test_awslambda/test_lambda.py
index 317e9f4a2..7bdfe3256 100644
--- a/tests/test_awslambda/test_lambda.py
+++ b/tests/test_awslambda/test_lambda.py
@@ -488,6 +488,7 @@ def lambda_handler(event, context):
     assert 'FunctionError' in result
     assert result['FunctionError'] == 'Handled'
 
+
 @mock_lambda
 @mock_s3
 def test_tags():
@@ -554,6 +555,7 @@ def test_tags():
         TagKeys=['spam']
     )['ResponseMetadata']['HTTPStatusCode'].should.equal(204)
 
+
 @mock_lambda
 def test_tags_not_found():
     """
@@ -574,6 +576,7 @@ def test_tags_not_found():
         TagKeys=['spam']
     ).should.throw(botocore.client.ClientError)
 
+
 @mock_lambda
 def test_invoke_async_function():
     conn = boto3.client('lambda', 'us-west-2')
@@ -581,10 +584,8 @@ def test_invoke_async_function():
         FunctionName='testFunction',
         Runtime='python2.7',
         Role='test-iam-role',
-        Handler='lambda_function.handler',
-        Code={
-            'ZipFile': get_test_zip_file1(),
-        },
+        Handler='lambda_function.lambda_handler',
+        Code={'ZipFile': get_test_zip_file1()},
         Description='test lambda function',
         Timeout=3,
         MemorySize=128,
@@ -593,11 +594,12 @@ def test_invoke_async_function():
 
     success_result = conn.invoke_async(
         FunctionName='testFunction', 
-        InvokeArgs=json.dumps({ 'test': 'event' })
+        InvokeArgs=json.dumps({'test': 'event'})
         )
 
     success_result['Status'].should.equal(202)
 
+
 @mock_lambda
 @freeze_time('2015-01-01 00:00:00')
 def test_get_function_created_with_zipfile():
@@ -646,6 +648,7 @@ def test_get_function_created_with_zipfile():
         },
     )
 
+
 @mock_lambda
 def add_function_permission():
     conn = boto3.client('lambda', 'us-west-2')