Support CustomAmiId in EMR (#2037)

moto/emr/models.py

@@ -97,7 +97,8 @@ class FakeCluster(BaseModel):
                  visible_to_all_users='false',
                  release_label=None,
                  requested_ami_version=None,
-                 running_ami_version=None):
+                 running_ami_version=None,
+                 custom_ami_id=None):
         self.id = cluster_id or random_cluster_id()
         emr_backend.clusters[self.id] = self
         self.emr_backend = emr_backend
@@ -162,6 +163,7 @@ class FakeCluster(BaseModel):
         self.release_label = release_label
         self.requested_ami_version = requested_ami_version
         self.running_ami_version = running_ami_version
+        self.custom_ami_id = custom_ami_id
 
         self.role = job_flow_role or 'EMRJobflowDefault'
         self.service_role = service_role

moto/emr/responses.py

@@ -267,6 +267,18 @@ class ElasticMapReduceResponse(BaseResponse):
             else:
                 kwargs['running_ami_version'] = '1.0.0'
 
+        custom_ami_id = self._get_param('CustomAmiId')
+        if custom_ami_id:
+            kwargs['custom_ami_id'] = custom_ami_id
+            if release_label and release_label < 'emr-5.7.0':
+                message = 'Custom AMI is not allowed'
+                raise EmrError(error_type='ValidationException',
+                            message=message, template='error_json')
+            elif ami_version:
+                message = 'Custom AMI is not supported in this version of EMR'
+                raise EmrError(error_type='ValidationException',
+                            message=message, template='error_json')
+
         cluster = self.backend.run_job_flow(**kwargs)
 
         applications = self._get_list_prefix('Applications.member')
@@ -375,6 +387,9 @@ DESCRIBE_CLUSTER_TEMPLATE = """<DescribeClusterResponse xmlns="http://elasticmap
         </member>
         {% endfor %}
       </Configurations>
+      {% if cluster.custom_ami_id is not none %}
+      <CustomAmiId>{{ cluster.custom_ami_id }}</CustomAmiId>
+      {% endif %}
       <Ec2InstanceAttributes>
         <AdditionalMasterSecurityGroups>
         {% for each in cluster.additional_master_security_groups %}

tests/test_emr/test_emr_boto3.py

@@ -432,6 +432,47 @@ def test_run_job_flow_with_instance_groups():
             x['BidPrice'].should.equal(y['BidPrice'])
 
 
+@mock_emr
+def test_run_job_flow_with_custom_ami():
+    client = boto3.client('emr', region_name='us-east-1')
+
+    with assert_raises(ClientError) as ex:
+        # CustomAmiId available in Amazon EMR 5.7.0 and later
+        args = deepcopy(run_job_flow_args)
+        args['CustomAmiId'] = 'MyEmrCustomId'
+        args['ReleaseLabel'] = 'emr-5.6.0'
+        client.run_job_flow(**args)
+    ex.exception.response['Error']['Code'].should.equal('ValidationException')
+    ex.exception.response['Error']['Message'].should.equal('Custom AMI is not allowed')
+
+    with assert_raises(ClientError) as ex:
+        args = deepcopy(run_job_flow_args)
+        args['CustomAmiId'] = 'MyEmrCustomId'
+        args['AmiVersion'] = '3.8.1'
+        client.run_job_flow(**args)
+    ex.exception.response['Error']['Code'].should.equal('ValidationException')
+    ex.exception.response['Error']['Message'].should.equal(
+        'Custom AMI is not supported in this version of EMR')
+
+    with assert_raises(ClientError) as ex:
+        # AMI version and release label exception  raises before CustomAmi exception
+        args = deepcopy(run_job_flow_args)
+        args['CustomAmiId'] = 'MyEmrCustomId'
+        args['ReleaseLabel'] = 'emr-5.6.0'
+        args['AmiVersion'] = '3.8.1'
+        client.run_job_flow(**args)
+    ex.exception.response['Error']['Code'].should.equal('ValidationException')
+    ex.exception.response['Error']['Message'].should.contain(
+        'Only one AMI version and release label may be specified.')
+
+    args = deepcopy(run_job_flow_args)
+    args['CustomAmiId'] = 'MyEmrCustomAmi'
+    args['ReleaseLabel'] = 'emr-5.7.0'
+    cluster_id = client.run_job_flow(**args)['JobFlowId']
+    resp = client.describe_cluster(ClusterId=cluster_id)
+    resp['Cluster']['CustomAmiId'].should.equal('MyEmrCustomAmi')
+
+
 @mock_emr
 def test_set_termination_protection():
     client = boto3.client('emr', region_name='us-east-1')