From dff24cb032fc11d80b50034423613b5739b9b492 Mon Sep 17 00:00:00 2001
From: Seth Black <seth.black@cleardata.com>
Date: Wed, 9 Oct 2019 16:20:49 -0500
Subject: [PATCH] bringing up test percentage

---
 moto/awslambda/models.py            |  3 --
 moto/awslambda/responses.py         |  3 ++
 tests/test_awslambda/test_lambda.py | 81 ++++++++++++++++++++++++++---
 3 files changed, 77 insertions(+), 10 deletions(-)

diff --git a/moto/awslambda/models.py b/moto/awslambda/models.py
index a1643cd03..bd85ded9a 100644
--- a/moto/awslambda/models.py
+++ b/moto/awslambda/models.py
@@ -296,9 +296,6 @@ class LambdaFunction(BaseModel):
         if 'DryRun' in updated_spec and updated_spec['DryRun']:
             return self.get_configuration()
 
-        if 'Publish' in updated_spec and updated_spec['Publish']:
-            self.set_version(self.version + 1)
-
         if 'ZipFile' in updated_spec:
             self.code['ZipFile'] = updated_spec['ZipFile']
 
diff --git a/moto/awslambda/responses.py b/moto/awslambda/responses.py
index 2041aa660..83a1cefca 100644
--- a/moto/awslambda/responses.py
+++ b/moto/awslambda/responses.py
@@ -342,6 +342,9 @@ class LambdaResponse(BaseResponse):
         fn = self.lambda_backend.get_function(function_name, qualifier)
 
         if fn:
+            if self.json_body.get('Publish', False):
+                fn = self.lambda_backend.publish_function(function_name)
+
             config = fn.update_function_code(self.json_body)
             return 200, {}, json.dumps(config)
         else:
diff --git a/tests/test_awslambda/test_lambda.py b/tests/test_awslambda/test_lambda.py
index 20a806de2..d6ee3f7f9 100644
--- a/tests/test_awslambda/test_lambda.py
+++ b/tests/test_awslambda/test_lambda.py
@@ -1295,13 +1295,13 @@ def test_update_configuration():
 
 
 @mock_lambda
-def test_update_function():
+def test_update_function_zip():
     conn = boto3.client('lambda', 'us-west-2')
 
     zip_content_one = get_test_zip_file1()
 
     fxn = conn.create_function(
-        FunctionName='testFunction',
+        FunctionName='testFunctionZip',
         Runtime='python2.7',
         Role='test-iam-role',
         Handler='lambda_function.lambda_handler',
@@ -1317,13 +1317,14 @@ def test_update_function():
     zip_content_two = get_test_zip_file2()
 
     fxn_updated = conn.update_function_code(
-        FunctionName='testFunction',
+        FunctionName='testFunctionZip',
         ZipFile=zip_content_two,
         Publish=True
     )
 
     response = conn.get_function(
-        FunctionName='testFunction'
+        FunctionName='testFunctionZip',
+        Qualifier='2'
     )
     response['Configuration'].pop('LastModified')
 
@@ -1336,14 +1337,80 @@ def test_update_function():
             "CodeSha256": hashlib.sha256(zip_content_two).hexdigest(),
             "CodeSize": len(zip_content_two),
             "Description": "test lambda function",
-            "FunctionArn": 'arn:aws:lambda:{}:123456789012:function:testFunction:2'.format(_lambda_region),
-            "FunctionName": "testFunction",
+            "FunctionArn": 'arn:aws:lambda:{}:123456789012:function:testFunctionZip:2'.format(_lambda_region),
+            "FunctionName": "testFunctionZip",
             "Handler": "lambda_function.lambda_handler",
             "MemorySize": 128,
             "Role": "test-iam-role",
             "Runtime": "python2.7",
             "Timeout": 3,
-            "Version": '$LATEST',
+            "Version": '2',
+            "VpcConfig": {
+                "SecurityGroupIds": [],
+                "SubnetIds": [],
+            }
+        },
+    )
+
+@mock_lambda
+@mock_s3
+def test_update_function_s3():
+    s3_conn = boto3.client('s3', 'us-west-2')
+    s3_conn.create_bucket(Bucket='test-bucket')
+
+    zip_content = get_test_zip_file1()
+    s3_conn.put_object(Bucket='test-bucket', Key='test.zip', Body=zip_content)
+
+    conn = boto3.client('lambda', 'us-west-2')
+
+    fxn = conn.create_function(
+        FunctionName='testFunctionS3',
+        Runtime='python2.7',
+        Role='test-iam-role',
+        Handler='lambda_function.lambda_handler',
+        Code={
+            'S3Bucket': 'test-bucket',
+            'S3Key': 'test.zip',
+        },
+        Description='test lambda function',
+        Timeout=3,
+        MemorySize=128,
+        Publish=True,
+    )
+
+    zip_content_two = get_test_zip_file2()
+    s3_conn.put_object(Bucket='test-bucket', Key='test2.zip', Body=zip_content_two)
+
+    fxn_updated = conn.update_function_code(
+        FunctionName='testFunctionS3',
+        S3Bucket='test-bucket',
+        S3Key='test2.zip',
+        Publish=True
+    )
+
+    response = conn.get_function(
+        FunctionName='testFunctionS3',
+        Qualifier='2'
+    )
+    response['Configuration'].pop('LastModified')
+
+    response['ResponseMetadata']['HTTPStatusCode'].should.equal(200)
+    assert len(response['Code']) == 2
+    assert response['Code']['RepositoryType'] == 'S3'
+    assert response['Code']['Location'].startswith('s3://awslambda-{0}-tasks.s3-{0}.amazonaws.com'.format(_lambda_region))
+    response['Configuration'].should.equal(
+        {
+            "CodeSha256": hashlib.sha256(zip_content_two).hexdigest(),
+            "CodeSize": len(zip_content_two),
+            "Description": "test lambda function",
+            "FunctionArn": 'arn:aws:lambda:{}:123456789012:function:testFunctionS3:2'.format(_lambda_region),
+            "FunctionName": "testFunctionS3",
+            "Handler": "lambda_function.lambda_handler",
+            "MemorySize": 128,
+            "Role": "test-iam-role",
+            "Runtime": "python2.7",
+            "Timeout": 3,
+            "Version": '2',
             "VpcConfig": {
                 "SecurityGroupIds": [],
                 "SubnetIds": [],