Fix existing test failures

Author: dipanjank

src/sagemaker/lambda_helper.py

@@ -153,7 +153,9 @@ def update(self):
             try:
                 if self.script is not None:
                     response = lambda_client.update_function_code(
-                        FunctionName=self.function_name, ZipFile=_zip_lambda_code(self.script)
+                        FunctionName=self.function_name,
+                        ZipFile=_zip_lambda_code(self.script),
+                        Architectures=self.architectures,
                     )
                 else:
                     response = lambda_client.update_function_code(
@@ -165,7 +167,7 @@ def update(self):
                             zipped_code_dir=self.zipped_code_dir,
                             s3_bucket=self.s3_bucket,
                         ),
-                        Architectures=self.architectures
+                        Architectures=self.architectures,
                     )
                 return response
             except ClientError as e:

tests/unit/test_lambda_helper.py

@@ -175,6 +175,9 @@ def test_create_lambda_happycase1(sagemaker_session):
         Code=code,
         Timeout=120,
         MemorySize=128,
+        Architectures=None,
+        VpcConfig=None,
+        Environment=None,
     )
 
 
@@ -200,6 +203,9 @@ def test_create_lambda_happycase2(sagemaker_session):
         Code=code,
         Timeout=120,
         MemorySize=128,
+        Architectures=None,
+        VpcConfig=None,
+        Environment=None,
     )
 
 
@@ -252,7 +258,7 @@ def test_update_lambda_happycase1(sagemaker_session):
     lambda_obj.update()
 
     sagemaker_session.lambda_client.update_function_code.assert_called_with(
-        FunctionName=FUNCTION_NAME, ZipFile=ZIPPED_CODE
+        FunctionName=FUNCTION_NAME, ZipFile=ZIPPED_CODE, Architectures=None,
     )
 
 
@@ -270,7 +276,7 @@ def test_update_lambda_happycase2(sagemaker_session):
     lambda_obj.update()
 
     sagemaker_session.lambda_client.update_function_code.assert_called_with(
-        FunctionName=LAMBDA_ARN, S3Bucket=S3_BUCKET, S3Key=S3_KEY
+        FunctionName=LAMBDA_ARN, S3Bucket=S3_BUCKET, S3Key=S3_KEY, Architectures=None,
     )
 
 
@@ -302,12 +308,6 @@ def test_upsert_lambda_happycase1(sagemaker_session):
         script=SCRIPT,
         handler=HANDLER,
         session=sagemaker_session,
-        architectures=["x86_64"],
-        environment={"Name": "my-test-lambda"},
-        vpc_config={
-            "SubnetIds": ["test-subnet-1"],
-            "SecurityGroupIds": ["sec-group-1"]
-        }
     )
 
     code = {"ZipFile": ZIPPED_CODE}
@@ -321,12 +321,9 @@ def test_upsert_lambda_happycase1(sagemaker_session):
         Code=code,
         Timeout=120,
         MemorySize=128,
-        Architectures=["x86_64"],
-        Environment={"Name": "my-test-lambda"},
-        VpcConfig={
-            "SubnetIds": ["test-subnet-1"],
-            "SecurityGroupIds": ["sec-group-1"]
-        }
+        Architectures=None,
+        VpcConfig=None,
+        Environment=None,
     )
 
 
@@ -338,7 +335,6 @@ def test_upsert_lambda_happycase2(sagemaker_session):
         script=SCRIPT,
         handler=HANDLER,
         session=sagemaker_session,
-        architectures=["x86_64"]
     )
 
     sagemaker_session.lambda_client.create_function.side_effect = ClientError(
@@ -349,7 +345,7 @@ def test_upsert_lambda_happycase2(sagemaker_session):
     lambda_obj.upsert()
 
     sagemaker_session.lambda_client.update_function_code.assert_called_once_with(
-        FunctionName=FUNCTION_NAME, ZipFile=ZIPPED_CODE, Architectures=["x86_64"]
+        FunctionName=FUNCTION_NAME, ZipFile=ZIPPED_CODE, Architectures=None
     )