diff --git a/moto/emr/models.py b/moto/emr/models.py index 466467360..5193a1986 100644 --- a/moto/emr/models.py +++ b/moto/emr/models.py @@ -124,6 +124,10 @@ class FakeJobFlow(object): else: self.visible_to_all_users = False + + def set_termination_protection(self, value): + self.termination_protected = value + def add_steps(self, steps): for index, step in enumerate(steps): if self.steps: @@ -236,6 +240,11 @@ class ElasticMapReduceBackend(BaseBackend): job = self.job_flows[job_id] job.set_visibility(visible_to_all_users) + def set_termination_protection(self, job_ids, value): + for job_id in job_ids: + job = self.job_flows[job_id] + job.set_termination_protection(value) + def add_tags(self, cluster_id, tags): cluster = self.get_cluster(cluster_id) cluster.add_tags(tags) diff --git a/moto/emr/responses.py b/moto/emr/responses.py index 337d0417b..926704c10 100644 --- a/moto/emr/responses.py +++ b/moto/emr/responses.py @@ -27,6 +27,10 @@ class ElasticMapReduceResponse(BaseResponse): flow_name, log_uri, job_flow_role, visible_to_all_users, steps, instance_attrs ) + instance_groups = self._get_list_prefix('Instances.InstanceGroups.member') + if instance_groups: + emr_backend.add_instance_groups(job_flow.id, instance_groups) + template = self.response_template(RUN_JOB_FLOW_TEMPLATE) return template.render(job_flow=job_flow) @@ -62,6 +66,13 @@ class ElasticMapReduceResponse(BaseResponse): template = self.response_template(SET_VISIBLE_TO_ALL_USERS_TEMPLATE) return template.render() + def set_termination_protection(self): + termination_protection = self._get_param('TerminationProtected') + job_ids = self._get_multi_param('JobFlowIds.member') + emr_backend.set_termination_protection(job_ids, termination_protection) + template = self.response_template(SET_TERMINATION_PROTECTION_TEMPLATE) + return template.render() + def list_clusters(self): clusters = emr_backend.list_clusters() template = self.response_template(LIST_CLUSTERS_TEMPLATE) @@ -149,6 +160,7 @@ DESCRIBE_JOB_FLOWS_TEMPLATE = """ + + + 2690d7eb-ed86-11dd-9877-6fad448a8419 + + +""" + ADD_TAGS_TEMPLATE = """ diff --git a/tests/test_emr/test_emr.py b/tests/test_emr/test_emr.py index 9157a82a5..6667f04db 100644 --- a/tests/test_emr/test_emr.py +++ b/tests/test_emr/test_emr.py @@ -1,6 +1,8 @@ from __future__ import unicode_literals + import boto from boto.emr.instance_group import InstanceGroup + from boto.emr.step import StreamingStep import sure # noqa @@ -107,6 +109,27 @@ def test_create_job_flow_visible_to_all_users(): job_flow.visibletoallusers.should.equal('True') +@requires_boto_gte("2.8") +@mock_emr +def test_create_job_flow_with_instance_groups(): + conn = boto.connect_emr() + + instance_groups = [InstanceGroup(6, 'TASK', 'c1.medium', 'SPOT', 'spot-0.07', '0.07'), + InstanceGroup(6, 'TASK', 'c1.medium', 'SPOT', 'spot-0.07', '0.07')] + job_id = conn.run_jobflow( + name='My jobflow', + log_uri='s3://some_bucket/jobflow_logs', + steps=[], + instance_groups=instance_groups + ) + + job_flow = conn.describe_jobflow(job_id) + int(job_flow.instancecount).should.equal(12) + instance_group = job_flow.instancegroups[0] + int(instance_group.instancerunningcount).should.equal(6) + + + @mock_emr def test_terminate_job_flow(): conn = boto.connect_emr() @@ -318,6 +341,30 @@ def test_set_visible_to_all_users(): job_flow.visibletoallusers.should.equal('False') +@requires_boto_gte("2.8") +@mock_emr +def test_set_termination_protection(): + conn = boto.connect_emr() + + job_id = conn.run_jobflow( + name='My jobflow', + log_uri='s3://some_bucket/jobflow_logs', + steps=[] + ) + job_flow = conn.describe_jobflow(job_id) + job_flow.terminationprotected.should.equal(u'None') + + conn.set_termination_protection(job_id, True) + + job_flow = conn.describe_jobflow(job_id) + job_flow.terminationprotected.should.equal('true') + + conn.set_termination_protection(job_id, False) + + job_flow = conn.describe_jobflow(job_id) + job_flow.terminationprotected.should.equal('false') + + @mock_emr def test_list_clusters(): conn = boto.connect_emr()