Merge pull request #27 from mathom/roleupdate

Update the IAM role if it exists

tests/placebo/TestZappa.test_create_iam_roles/iam.CreateRole_1.json

@@ -0,0 +1,26 @@
+{
+    "status_code": 200, 
+    "data": {
+        "Role": {
+            "AssumeRolePolicyDocument": "%7B%0A%20%20%22Version%22%3A%20%222012-10-17%22%2C%0A%20%20%22Statement%22%3A%20%5B%0A%20%20%20%20%7B%0A%20%20%20%20%20%20%22Sid%22%3A%20%22%22%2C%0A%20%20%20%20%20%20%22Effect%22%3A%20%22Allow%22%2C%0A%20%20%20%20%20%20%22Principal%22%3A%20%7B%0A%20%20%20%20%20%20%20%20%22Service%22%3A%20%5B%0A%20%20%20%20%20%20%20%20%20%20%22lambda.amazonaws.com%22%2C%0A%20%20%20%20%20%20%20%20%20%20%22apigateway.amazonaws.com%22%0A%20%20%20%20%20%20%20%20%5D%0A%20%20%20%20%20%20%7D%2C%0A%20%20%20%20%20%20%22Action%22%3A%20%22sts%3AAssumeRole%22%0A%20%20%20%20%7D%0A%20%20%5D%0A%7D", 
+            "RoleId": "AROAIJTKAN7OBBD7BF25O", 
+            "CreateDate": {
+                "hour": 1, 
+                "__class__": "datetime", 
+                "month": 2, 
+                "second": 20, 
+                "microsecond": 221000, 
+                "year": 2016, 
+                "day": 25, 
+                "minute": 28
+            }, 
+            "RoleName": "ZappaLambdaExecution", 
+            "Path": "/", 
+            "Arn": "arn:aws:iam::123:role/ZappaLambdaExecution"
+        }, 
+        "ResponseMetadata": {
+            "HTTPStatusCode": 200, 
+            "RequestId": "0da2623f-db5f-11e5-9d89-67e11bce6785"
+        }
+    }
+}

tests/placebo/TestZappa.test_create_iam_roles/iam.GetRolePolicy_1.json

@@ -0,0 +1,14 @@
+{
+    "status_code": 404, 
+    "data": {
+        "ResponseMetadata": {
+            "HTTPStatusCode": 404, 
+            "RequestId": "0e11c69f-db5f-11e5-aa20-65c7ceb91dfa"
+        }, 
+        "Error": {
+            "Message": "The role policy with name zappa-permissions cannot be found.", 
+            "Code": "NoSuchEntity", 
+            "Type": "Sender"
+        }
+    }
+}

tests/placebo/TestZappa.test_create_iam_roles/iam.GetRole_1.json

@@ -0,0 +1,14 @@
+{
+    "status_code": 404, 
+    "data": {
+        "ResponseMetadata": {
+            "HTTPStatusCode": 404, 
+            "RequestId": "0d24322f-db5f-11e5-9d89-67e11bce6785"
+        }, 
+        "Error": {
+            "Message": "The role with name ZappaLambdaExecution cannot be found.", 
+            "Code": "NoSuchEntity", 
+            "Type": "Sender"
+        }
+    }
+}

tests/placebo/TestZappa.test_create_iam_roles/iam.PutRolePolicy_1.json

@@ -0,0 +1,9 @@
+{
+    "status_code": 200, 
+    "data": {
+        "ResponseMetadata": {
+            "HTTPStatusCode": 200, 
+            "RequestId": "0e619422-db5f-11e5-9d89-67e11bce6785"
+        }
+    }
+}

tests/tests.py

@@ -1,21 +1,15 @@
-import boto3
 import collections
+import json
 import os
-import placebo
 import unittest
 
+from .utils import placebo_session
+
 from zappa.wsgi import create_wsgi_request, common_log
-from zappa.zappa import Zappa
+from zappa.zappa import Zappa, ASSUME_POLICY, ATTACH_POLICY
 
 
 class TestZappa(unittest.TestCase):
-    def get_placebo_session(self):
-        session = boto3.Session()
-        placebo_dir = os.path.join(os.path.dirname(__file__), 'placebo')
-        pill = placebo.attach(session, data_path=placebo_dir)
-        pill.playback()
-        return session
-
     ##
     # Sanity Tests
     ##
@@ -60,8 +54,8 @@ def test_load_credentials(self):
         self.assertTrue((z.secret_key == "JKL456"))
         self.assertTrue((z.aws_region == 'us-east-1'))
 
-    def test_upload_remove_s3(self):
-        session = self.get_placebo_session()
+    @placebo_session
+    def test_upload_remove_s3(self, session):
         bucket_name = 'test_zappa_upload_s3'
         z = Zappa()
         zip_path = z.create_lambda_zip()
@@ -81,6 +75,17 @@ def test_upload_remove_s3(self):
         res = z.remove_from_s3(zip_path, bucket_name, session)
         self.assertTrue(res)
 
+    @placebo_session
+    def test_create_iam_roles(self, session):
+        z = Zappa()
+        arn = z.create_iam_roles(session)
+        self.assertEqual(arn, "arn:aws:iam::123:role/{}".format(z.role_name))
+
+    def test_policy_json(self):
+        # ensure the policy docs are valid JSON
+        json.loads(ASSUME_POLICY)
+        json.loads(ATTACH_POLICY)
+
     ##
     # Logging
     ##
@@ -89,7 +94,7 @@ def test_logging(self):
         """
         TODO
         """
-        z = Zappa()
+        Zappa()
 
     ##
     # WSGI

tests/utils.py

@@ -0,0 +1,49 @@
+import placebo
+import boto3
+import os
+import functools
+
+PLACEBO_DIR = os.path.join(os.path.dirname(__file__), 'placebo')
+
+
+def placebo_session(function):
+    """
+    Decorator to help do testing with placebo.
+
+    Simply wrap the function you want to test and make sure to add
+    a "session" argument so the decorator can pass the placebo session.
+
+    Accepts the following environment variables to configure placebo:
+
+    PLACEBO_MODE: set to "record" to record AWS calls and save them
+    PLACEBO_PROFILE: optionally set an AWS credential profile to record with
+    """
+
+    @functools.wraps(function)
+    def wrapper(*args, **kwargs):
+        session_kwargs = {}
+        profile_name = os.environ.get('PLACEBO_PROFILE', None)
+        if profile_name:
+            session_kwargs['profile_name'] = profile_name
+
+        session = boto3.Session(**session_kwargs)
+
+        self = args[0]
+        prefix = self.__class__.__name__ + '.' + function.__name__
+        record_dir = os.path.join(PLACEBO_DIR, prefix)
+
+        if not os.path.exists(record_dir):
+            os.makedirs(record_dir)
+
+        pill = placebo.attach(session, data_path=record_dir)
+
+        if os.environ.get('PLACEBO_MODE') == 'record':
+            pill.record()
+        else:
+            pill.playback()
+
+        kwargs['session'] = session
+
+        return function(*args, **kwargs)
+
+    return wrapper

zappa/zappa.py

@@ -6,6 +6,7 @@
 import time
 import zipfile
 import requests
+import json
 
 from os.path import expanduser
 from tqdm import tqdm
@@ -652,33 +653,43 @@ def get_api_url(self, stage_name):
     # IAM
     ##
 
-    def create_iam_roles(self):
+    def create_iam_roles(self, session=None):
         """
         Creates and defines the IAM roles and policies necessary for Zappa.
 
+        If the IAM role already exists, it will be updated if necessary.
         """
-
         assume_policy_s = ASSUME_POLICY
         attach_policy_s = ATTACH_POLICY
 
-        iam = boto3.resource('iam')
+        attach_policy_obj = json.loads(attach_policy_s)
+
+        session = session or boto3.session.Session()
+        iam = session.resource('iam')
 
+        # create the role if needed
+        role = iam.Role(self.role_name)
         try:
-            role = iam.meta.client.get_role(
-                RoleName=self.role_name)
-            self.credentials_arn = role['Role']['Arn']
-            return self.credentials_arn
+            self.credentials_arn = role.arn
 
         except botocore.client.ClientError:
-            print("Creating " + self.role_name + " IAM..")
+            print("Creating " + self.role_name + " IAM Role...")
+
+            role = iam.create_role(RoleName=self.role_name,
+                                   AssumeRolePolicyDocument=assume_policy_s)
+            self.credentials_arn = role.arn
 
-            role = iam.create_role(
-                RoleName=self.role_name,
-                AssumeRolePolicyDocument=assume_policy_s)
-            iam.RolePolicy(self.role_name, 'zappa-permissions').put(
-                PolicyDocument=attach_policy_s)
+        # create or update the role's policy if needed
+        policy = iam.RolePolicy(self.role_name, 'zappa-permissions')
+        try:
+            if policy.policy_document != attach_policy_obj:
+                print("Updating zappa-permissions policy on " + self.role_name + " IAM Role.")
+                policy.put(PolicyDocument=attach_policy_s)
+
+        except botocore.client.ClientError:
+            print("Creating zappa-permissions policy on " + self.role_name + " IAM Role.")
+            policy.put(PolicyDocument=attach_policy_s)
 
-        self.credentials_arn = role.arn
         return self.credentials_arn
 
     ##