From e05a0619939cabaa8b3b1c7054414db0f635c53e Mon Sep 17 00:00:00 2001
From: Steve Pulec <spulec@gmail.com>
Date: Wed, 7 Jan 2015 22:32:51 -0500
Subject: [PATCH] Add support for Cloudformation Mappings.

---
 moto/cloudformation/parsing.py                | 28 ++++++++----
 .../test_cloudformation_stack_integration.py  | 45 +++++++++++++++++++
 2 files changed, 65 insertions(+), 8 deletions(-)

diff --git a/moto/cloudformation/parsing.py b/moto/cloudformation/parsing.py
index 48881a0a4..8650b4f1b 100644
--- a/moto/cloudformation/parsing.py
+++ b/moto/cloudformation/parsing.py
@@ -73,6 +73,14 @@ def clean_json(resource_json, resources_map):
             else:
                 return resource
 
+        if "Fn::FindInMap" in resource_json:
+            map_name = resource_json["Fn::FindInMap"][0]
+            map_path = resource_json["Fn::FindInMap"][1:]
+            result = resources_map[map_name]
+            for path in map_path:
+                result = result[clean_json(path, resources_map)]
+            return result
+
         if 'Fn::GetAtt' in resource_json:
             resource = resources_map[resource_json['Fn::GetAtt'][0]]
             if resource is None:
@@ -88,13 +96,6 @@ def clean_json(resource_json, resources_map):
                     UnformattedGetAttTemplateException.description.format(
                         resource_json['Fn::GetAtt'][0], resource_json['Fn::GetAtt'][1]))
 
-        if 'Fn::Join' in resource_json:
-            join_list = []
-            for val in resource_json['Fn::Join'][1]:
-                cleaned_val = clean_json(val, resources_map)
-                join_list.append(cleaned_val if cleaned_val else '{0}'.format(val))
-            return resource_json['Fn::Join'][0].join(join_list)
-
         if 'Fn::If' in resource_json:
             condition_name, true_value, false_value = resource_json['Fn::If']
             if resources_map[condition_name]:
@@ -102,6 +103,13 @@ def clean_json(resource_json, resources_map):
             else:
                 return false_value
 
+        if 'Fn::Join' in resource_json:
+            join_list = []
+            for val in resource_json['Fn::Join'][1]:
+                cleaned_val = clean_json(val, resources_map)
+                join_list.append(cleaned_val if cleaned_val else '{0}'.format(val))
+            return resource_json['Fn::Join'][0].join(join_list)
+
         cleaned_json = {}
         for key, value in resource_json.items():
             cleaned_json[key] = clean_json(value, resources_map)
@@ -216,7 +224,7 @@ class ResourceMap(collections.Mapping):
         # Create the default resources
         self._parsed_resources = {
             "AWS::AccountId": "123456789012",
-            "AWS::Region": "us-east-1",
+            "AWS::Region": self._region_name,
             "AWS::StackId": stack_id,
             "AWS::StackName": stack_name,
         }
@@ -242,6 +250,9 @@ class ResourceMap(collections.Mapping):
     def resources(self):
         return self._resource_json_map.keys()
 
+    def load_mapping(self):
+        self._parsed_resources.update(self._template.get('Mappings', {}))
+
     def load_parameters(self):
         parameter_slots = self._template.get('Parameters', {})
         for parameter_name, parameter in parameter_slots.items():
@@ -267,6 +278,7 @@ class ResourceMap(collections.Mapping):
             self._parsed_resources[condition_name] = parse_condition(condition, self._parsed_resources)
 
     def create(self):
+        self.load_mapping()
         self.load_parameters()
         self.load_conditions()
 
diff --git a/tests/test_cloudformation/test_cloudformation_stack_integration.py b/tests/test_cloudformation/test_cloudformation_stack_integration.py
index df9a21584..b88be5e70 100644
--- a/tests/test_cloudformation/test_cloudformation_stack_integration.py
+++ b/tests/test_cloudformation/test_cloudformation_stack_integration.py
@@ -652,3 +652,48 @@ def test_conditional_if_handling():
     reservation = ec2_conn.get_all_instances()[0]
     ec2_instance = reservation.instances[0]
     ec2_instance.image_id.should.equal("ami-00000000")
+
+
+@mock_cloudformation()
+@mock_ec2()
+def test_cloudformation_mapping():
+    dummy_template = {
+        "AWSTemplateFormatVersion": "2010-09-09",
+        "Mappings": {
+            "RegionMap": {
+                "us-east-1": {"32": "ami-6411e20d", "64": "ami-7a11e213"},
+                "us-west-1": {"32": "ami-c9c7978c", "64": "ami-cfc7978a"},
+                "eu-west-1": {"32": "ami-37c2f643", "64": "ami-31c2f645"},
+                "ap-southeast-1": {"32": "ami-66f28c34", "64": "ami-60f28c32"},
+                "ap-northeast-1": {"32": "ami-9c03a89d", "64": "ami-a003a8a1"}
+            }
+        },
+        "Resources": {
+            "WebServer": {
+                "Type": "AWS::EC2::Instance",
+                "Properties": {
+                    "ImageId": {
+                        "Fn::FindInMap": ["RegionMap", {"Ref": "AWS::Region"}, "32"]
+                    },
+                    "InstanceType": "m1.small"
+                },
+                "Type": "AWS::EC2::Instance",
+            },
+        },
+    }
+
+    dummy_template_json = json.dumps(dummy_template)
+
+    conn = boto.cloudformation.connect_to_region("us-east-1")
+    conn.create_stack('test_stack1', template_body=dummy_template_json)
+    ec2_conn = boto.ec2.connect_to_region("us-east-1")
+    reservation = ec2_conn.get_all_instances()[0]
+    ec2_instance = reservation.instances[0]
+    ec2_instance.image_id.should.equal("ami-6411e20d")
+
+    conn = boto.cloudformation.connect_to_region("us-west-1")
+    conn.create_stack('test_stack1', template_body=dummy_template_json)
+    ec2_conn = boto.ec2.connect_to_region("us-west-1")
+    reservation = ec2_conn.get_all_instances()[0]
+    ec2_instance = reservation.instances[0]
+    ec2_instance.image_id.should.equal("ami-c9c7978c")