def test_can_inject_policy(self, sample_websocket_app):
config = Config.create(chalice_app=sample_websocket_app,
autogen_policy=True,
project_dir='.')
event_source = self.create_model_from_app(sample_websocket_app, config)
role = event_source.connect_function.role
role.policy.document = {'Statement': []}
injector = WebsocketPolicyInjector()
injector.handle(config, event_source)
assert role.policy.document == {
'Statement': [POST_TO_WEBSOCKET_CONNECTION_POLICY.copy()],
}
def _inject_into_function(self, config, lambda_function):
# type: (Config, Optional[models.LambdaFunction]) -> None
if lambda_function is None:
return
role = lambda_function.role
if role is None:
return
if (not self._policy_injected
and isinstance(role, models.ManagedIAMRole)
and isinstance(role.policy, models.AutoGenIAMPolicy)
and not isinstance(role.policy.document, models.Placeholder)):
self._inject_policy(role.policy.document,
POST_TO_WEBSOCKET_CONNECTION_POLICY.copy())
self._policy_injected = True