Fix issue in updating iam role arn for deployments

CHANGELOG.rst

@@ -2,6 +2,14 @@
 CHANGELOG
 =========
 
+Next Release (TBD)
+==================
+
+* Fix issue where provided ``iam_role_arn`` was not respected on
+  redeployments of chalice applications and in the CloudFormation template
+  generated by ``chalice package``
+  (`#339 <https://github.com/awslabs/chalice/issues/339>`__)
+
 0.9.0
 =====
 

chalice/awsclient.py

@@ -30,6 +30,7 @@
 _STR_MAP = Optional[Dict[str, str]]
 _OPT_STR = Optional[str]
 _OPT_INT = Optional[int]
+_CLIENT_METHOD = Callable[..., Dict[str, Any]]
 
 
 class ResourceDoesNotExistError(Exception):
@@ -90,22 +91,27 @@ def create_function(self,
             kwargs['Timeout'] = timeout
         if memory_size is not None:
             kwargs['MemorySize'] = memory_size
+        return self._call_client_method_with_retries(
+            self._client('lambda').create_function, kwargs)['FunctionArn']
+
+    def _call_client_method_with_retries(self, method, kwargs):
+        # type: (_CLIENT_METHOD, Dict[str, Any]) -> Dict[str, Any]
         client = self._client('lambda')
         attempts = 0
         while True:
             try:
-                response = client.create_function(**kwargs)
+                response = method(**kwargs)
             except client.exceptions.InvalidParameterValueException:
                 # We're assuming that if we receive an
                 # InvalidParameterValueException, it's because
                 # the role we just created can't be used by
-                # Lambda.
+                # Lambda so retry until it can be.
                 self._sleep(self.DELAY_TIME)
                 attempts += 1
                 if attempts >= self.LAMBDA_CREATE_ATTEMPTS:
                     raise
                 continue
-            return response['FunctionArn']
+            return response
 
     def delete_function(self, function_name):
         # type: (str) -> None
@@ -122,7 +128,8 @@ def update_function(self,
                         runtime=None,                # type: _OPT_STR
                         tags=None,                   # type: _STR_MAP
                         timeout=None,                # type: _OPT_INT
-                        memory_size=None             # type: _OPT_INT
+                        memory_size=None,            # type: _OPT_INT
+                        role_arn=None                # type: _OPT_STR
                         ):
         # type: (...) -> Dict[str, Any]
         """Update a Lambda function's code and configuration.
@@ -144,9 +151,12 @@ def update_function(self,
             kwargs['Timeout'] = timeout
         if memory_size is not None:
             kwargs['MemorySize'] = memory_size
+        if role_arn is not None:
+            kwargs['Role'] = role_arn
         if kwargs:
             kwargs['FunctionName'] = function_name
-            lambda_client.update_function_configuration(**kwargs)
+            self._call_client_method_with_retries(
+                lambda_client.update_function_configuration, kwargs)
         if tags is not None:
             self._update_function_tags(return_value['FunctionArn'], tags)
         return return_value

chalice/deploy/deployer.py

@@ -383,6 +383,7 @@ def _update_lambda_function(self, config, lambda_name, stage_name):
                 project_dir)
         zip_contents = self._osutils.get_file_contents(
             deployment_package_filename, binary=True)
+        role_arn = self._get_or_create_lambda_role_arn(config, lambda_name)
         print("Sending changes to lambda.")
         self._aws_client.update_function(
             function_name=lambda_name,
@@ -391,7 +392,8 @@ def _update_lambda_function(self, config, lambda_name, stage_name):
             environment_variables=config.environment_variables,
             tags=config.tags,
             timeout=self._get_lambda_timeout(config),
-            memory_size=self._get_lambda_memory_size(config)
+            memory_size=self._get_lambda_memory_size(config),
+            role_arn=role_arn
         )
 
     def _write_config_to_disk(self, config):

chalice/package.py

@@ -104,7 +104,6 @@ def _generate_serverless_function(self, config, code_uri):
             'Handler': 'app.app',
             'CodeUri': code_uri,
             'Events': self._generate_function_events(config.chalice_app),
-            'Policies': [self._generate_iam_policy()],
             'Tags': config.tags,
             'Timeout': DEFAULT_LAMBDA_TIMEOUT,
             'MemorySize': DEFAULT_LAMBDA_MEMORY_SIZE
@@ -117,6 +116,10 @@ def _generate_serverless_function(self, config, code_uri):
             properties['Timeout'] = config.lambda_timeout
         if config.lambda_memory_size is not None:
             properties['MemorySize'] = config.lambda_memory_size
+        if not config.manage_iam_role:
+            properties['Role'] = config.iam_role_arn
+        else:
+            properties['Policies'] = [self._generate_iam_policy()]
         return {
             'Type': 'AWS::Serverless::Function',
             'Properties': properties,

tests/functional/test_awsclient.py

@@ -595,6 +595,66 @@ def test_update_function_with_no_tag_updates_needed(self, stubbed_session):
         awsclient.update_function('name', b'foo', tags={'MyKey': 'SameValue'})
         stubbed_session.verify_stubs()
 
+    def test_update_function_with_iam_role(self, stubbed_session):
+        function_arn = 'arn'
+
+        lambda_client = stubbed_session.stub('lambda')
+        lambda_client.update_function_code(
+            FunctionName='name', ZipFile=b'foo').returns(
+                {'FunctionArn': function_arn})
+        lambda_client.update_function_configuration(
+            FunctionName='name',
+            Role='role-arn').returns({})
+        stubbed_session.activate_stubs()
+
+        awsclient = TypedAWSClient(stubbed_session)
+        awsclient.update_function('name', b'foo', role_arn='role-arn')
+        stubbed_session.verify_stubs()
+
+    def test_update_function_is_retried_and_succeeds(self, stubbed_session):
+        stubbed_session.stub('lambda').update_function_code(
+            FunctionName='name', ZipFile=b'foo').returns(
+                {'FunctionArn': 'arn'})
+
+        update_config_kwargs = {
+            'FunctionName': 'name',
+            'Role': 'role-arn'
+        }
+        # This should fail two times with retryable exceptions and
+        # then succeed to update the lambda function.
+        stubbed_session.stub('lambda').update_function_configuration(
+            **update_config_kwargs).raises_error(
+                error_code='InvalidParameterValueException', message='')
+        stubbed_session.stub('lambda').update_function_configuration(
+            **update_config_kwargs).raises_error(
+            error_code='InvalidParameterValueException', message='')
+        stubbed_session.stub('lambda').update_function_configuration(
+            **update_config_kwargs).returns({})
+        stubbed_session.activate_stubs()
+        awsclient = TypedAWSClient(stubbed_session, mock.Mock(spec=time.sleep))
+        awsclient.update_function('name', b'foo', role_arn='role-arn')
+        stubbed_session.verify_stubs()
+
+    def test_update_function_fails_after_max_retries(self, stubbed_session):
+        stubbed_session.stub('lambda').update_function_code(
+            FunctionName='name', ZipFile=b'foo').returns(
+                {'FunctionArn': 'arn'})
+
+        update_config_kwargs = {
+            'FunctionName': 'name',
+            'Role': 'role-arn'
+        }
+        for _ in range(TypedAWSClient.LAMBDA_CREATE_ATTEMPTS):
+            stubbed_session.stub('lambda').update_function_configuration(
+                **update_config_kwargs).raises_error(
+                    error_code='InvalidParameterValueException', message='')
+        stubbed_session.activate_stubs()
+        awsclient = TypedAWSClient(stubbed_session, mock.Mock(spec=time.sleep))
+
+        with pytest.raises(botocore.exceptions.ClientError):
+            awsclient.update_function('name', b'foo', role_arn='role-arn')
+        stubbed_session.verify_stubs()
+
 
 class TestCanDeleteRolePolicy(object):
     def test_can_delete_role_policy(self, stubbed_session):

tests/unit/deploy/test_deployer.py

@@ -452,7 +452,7 @@ def test_lambda_deployer_repeated_deploy(app_policy, sample_app):
         chalice_app=sample_app,
         manage_iam_role=False,
         app_name='appname',
-        iam_role_arn=True,
+        iam_role_arn='role-arn',
         project_dir='./myproject',
         environment_variables={"FOO": "BAR"},
         lambda_timeout=120,
@@ -488,7 +488,8 @@ def test_lambda_deployer_repeated_deploy(app_policy, sample_app):
                 chalice_version, 'dev', 'appname'),
             'mykey': 'myvalue'
         },
-        timeout=120, memory_size=256
+        timeout=120, memory_size=256,
+        role_arn='role-arn'
     )
 
 
@@ -833,7 +834,8 @@ def test_lambda_deployer_defaults(self, sample_app):
                     chalice_version),
             },
             environment_variables={},
-            timeout=60, memory_size=128
+            timeout=60, memory_size=128,
+            role_arn='role-arn'
         )
 
     def test_lambda_deployer_with_environment_vars(self, sample_app):
@@ -856,7 +858,8 @@ def test_lambda_deployer_with_environment_vars(self, sample_app):
                     chalice_version),
             },
             environment_variables={'FOO': 'BAR'},
-            timeout=60, memory_size=128
+            timeout=60, memory_size=128,
+            role_arn='role-arn'
         )
 
     def test_lambda_deployer_with_timeout_configured(self, sample_app):
@@ -879,7 +882,8 @@ def test_lambda_deployer_with_timeout_configured(self, sample_app):
                     chalice_version),
             },
             environment_variables={},
-            timeout=120, memory_size=128
+            timeout=120, memory_size=128,
+            role_arn='role-arn'
         )
 
     def test_lambda_deployer_with_memory_size_configured(self, sample_app):
@@ -902,7 +906,8 @@ def test_lambda_deployer_with_memory_size_configured(self, sample_app):
                     chalice_version),
             },
             environment_variables={},
-            timeout=60, memory_size=256
+            timeout=60, memory_size=256,
+            role_arn='role-arn'
         )
 
     def test_lambda_deployer_with_tags(self, sample_app):
@@ -926,5 +931,6 @@ def test_lambda_deployer_with_tags(self, sample_app):
                 'mykey': 'myvalue'
             },
             environment_variables={},
-            timeout=60, memory_size=128
+            timeout=60, memory_size=128,
+            role_arn='role-arn'
         )

tests/unit/test_package.py

@@ -76,6 +76,7 @@ def test_sam_injects_policy(sample_app,
     assert template['Resources']['APIHandler']['Properties']['Policies'] == [{
         'iam': 'policy',
     }]
+    assert 'Role' not in template['Resources']['APIHandler']['Properties']
 
 
 def test_sam_injects_swagger_doc(sample_app,
@@ -240,3 +241,20 @@ def test_maps_python_version(sample_app,
     expected = config.lambda_python_version
     actual = template['Resources']['APIHandler']['Properties']['Runtime']
     assert actual == expected
+
+
+def test_role_arn_added_to_function(sample_app,
+                                    mock_swagger_generator,
+                                    mock_policy_generator):
+    p = package.SAMTemplateGenerator(
+        mock_swagger_generator, mock_policy_generator)
+    mock_swagger_generator.generate_swagger.return_value = {
+        'swagger': 'document'
+    }
+    config = Config.create(
+        chalice_app=sample_app, api_gateway_stage='dev', app_name='myapp',
+        manage_iam_role=False, iam_role_arn='role-arn')
+    template = p.generate_sam_template(config)
+    properties = template['Resources']['APIHandler']['Properties']
+    assert properties['Role'] == 'role-arn'
+    assert 'Policies' not in properties