def test_get_two_task_instance_groups(self, mock_emr):
mock_instance_groups = mock_emr.return_value.list_instance_groups
mock_instance_groups.return_value = {
"InstanceGroups": [{
"InstanceGroupType": "MASTER",
"InstanceGroupName": "MyMasterGroup"
}, {
"InstanceGroupType": "CORE",
"InstanceGroupName": "MyCoreGroup"
}, {
"InstanceGroupType": "TASK",
"InstanceGroupName": "MyTaskGroup",
"BidPrice": 1.2
}, {
"InstanceGroupType": "TASK",
"InstanceGroupName": "MyOtherTaskGroup",
"BidPrice": 1.2
}]
}
groups = Emr(job_flow_id=self.job_flow,
region="eu-west-1").get_task_instance_groups()
self.assertListEqual(groups, [{
"InstanceGroupType": "TASK",
"InstanceGroupName": "MyTaskGroup",
"BidPrice": 1.2
}, {
"InstanceGroupType": "TASK",
"InstanceGroupName": "MyOtherTaskGroup",
"BidPrice": 1.2
}])
mock_instance_groups.assert_called_with(ClusterId=self.job_flow)
def lambda_handler(event, context):
job_flow_id = event["JobFlowId"]
threshold = float(event["Threshold"])
min_instances = int(event["MinInstances"])
max_instances = int(event["MaxInstances"])
office_hours_start = int(event["OfficeHoursStart"])
office_hours_end = int(event["OfficeHoursEnd"])
shutdown_time = int(event["ShutdownTime"])
parent_stack_id = event[
"ParentStackId"] if "ParentStackId" in event else None
stack_deletion_role = event[
"StackDeletionRole"] if "StackDeletionRole" in event else None
scaler = EmrScaler(emr=Emr(job_flow_id=job_flow_id,
min_instances=min_instances,
max_instances=max_instances),
min_instances=min_instances,
max_instances=max_instances,
office_hours_start=office_hours_start,
office_hours_end=office_hours_end,
shutdown_time=shutdown_time,
parent_stack=parent_stack_id,
stack_deletion_role=stack_deletion_role)
scaler.maybe_shutdown()
scaler.maybe_scale(threshold)
def test_scales_only_once(self, mock_get_task_instance_group, mock_emr):
mock_get_task_instance_group.return_value = [{
"Id":
self.emr_task_instance_group_id,
"RequestedInstanceCount":
5,
"BidPrice":
1.2,
"Name":
self.emr_task_instance_group_name,
"InstanceType":
self.emr_task_instance_group_instance_type
}, {
"Id":
self.emr_task_instance_group_id + "_expensive",
"RequestedInstanceCount":
5,
"BidPrice":
1.5,
"Name":
self.emr_task_instance_group_name,
"InstanceType":
self.emr_task_instance_group_instance_type
}]
mock_modify_instance_groups = mock_emr.return_value.modify_instance_groups
Emr(job_flow_id=self.job_flow,
min_instances=5,
max_instances=10,
region="eu-west-1").scale(direction=1)
mock_modify_instance_groups.assert_called_once_with(InstanceGroups=[{
"InstanceGroupId":
self.emr_task_instance_group_id + "_expensive",
"InstanceCount":
6
}])
def test_no_scaling_in_progress(self, mock_get_task_instance_group):
mock_get_task_instance_group.return_value = [{
"RequestedInstanceCount": 1,
"RunningInstanceCount": 1
}]
self.assertFalse(
Emr(job_flow_id=self.job_flow,
region="eu-west-1").scaling_in_progress())
def test_is_termination_protected_False(self, mock_emr):
mock_describe_cluster = mock_emr.return_value.describe_cluster
mock_describe_cluster.return_value = {
"Cluster": {
"TerminationProtected": False
}
}
self.assertFalse(
Emr(job_flow_id=self.job_flow).is_termination_protected())
def test_get_no_task_instance_group(self, mock_emr):
mock_instance_groups = mock_emr.return_value.list_instance_groups
mock_instance_groups.return_value = {
"InstanceGroups": [{
"InstanceGroupType": "MASTER",
"InstanceGroupName": "MyMasterGroup"
}, {
"InstanceGroupType": "CORE",
"InstanceGroupName": "MyCoreGroup"
}]
}
groups = Emr(job_flow_id=self.job_flow,
region="eu-west-1").get_task_instance_groups()
self.assertListEqual(groups, [])
def test_gets_pending_containers(self, mock_cw):
mock_stats = mock_cw.return_value.get_metric_statistics
mock_stats.return_value = {"Datapoints": [{"Maximum": 42}]}
max = Emr(job_flow_id=self.job_flow,
region="eu-west-1").get_pending_containers()
now = datetime.utcnow().replace(second=0, microsecond=0)
self.assertEqual(max, 42)
mock_stats.assert_called_with(
Namespace="AWS/ElasticMapReduce",
MetricName=self.cw_metric_container_pending,
StartTime=now - timedelta(minutes=5),
EndTime=now,
Period=300,
Statistics=[
"Maximum",
],
Unit="Count",
Dimensions=[{
"Name": "JobFlowId",
"Value": self.job_flow
}])
def test_no_scaling_because_below_lower_bound(self,
mock_get_task_instance_group,
mock_emr):
mock_get_task_instance_group.return_value = [{
"Id":
self.emr_task_instance_group_id,
"RequestedInstanceCount":
5,
"BidPrice":
1.2,
"Name":
self.emr_task_instance_group_name,
"InstanceType":
self.emr_task_instance_group_instance_type
}]
mock_modify_instance_groups = mock_emr.return_value.modify_instance_groups
Emr(job_flow_id=self.job_flow,
min_instances=5,
max_instances=10,
region="eu-west-1").scale(direction=-1)
mock_modify_instance_groups.assert_not_called()
def test_returns_average_of_last_hour(self, mock_cw):
mock_stats = mock_cw.return_value.get_metric_statistics
mock_stats.return_value = {"Datapoints": [{"Average": 42.0}]}
avg = Emr(job_flow_id=self.job_flow,
region="eu-west-1").get_average_of_last_hour(
self.cw_metric_avg)
now = datetime.utcnow().replace(second=0, microsecond=0)
self.assertEqual(avg, 42.0)
mock_stats.assert_called_with(Namespace="AWS/ElasticMapReduce",
MetricName=self.cw_metric_avg,
StartTime=now - timedelta(hours=1),
EndTime=now,
Period=3600,
Statistics=[
"Average",
],
Unit="Count",
Dimensions=[{
"Name": "JobFlowId",
"Value": self.job_flow
}])
def test_scaling_down_from_one_instance(self, mock_get_task_instance_group,
mock_emr):
mock_get_task_instance_group.return_value = [{
"Id":
self.emr_task_instance_group_id,
"RequestedInstanceCount":
1,
"BidPrice":
1.2,
"Name":
self.emr_task_instance_group_name,
"InstanceType":
self.emr_task_instance_group_instance_type
}]
mock_modify_instance_groups = mock_emr.return_value.modify_instance_groups
Emr(job_flow_id=self.job_flow,
min_instances=0,
max_instances=10,
region="eu-west-1").scale(direction=-1)
mock_modify_instance_groups.assert_called_with(
InstanceGroups=[{
"InstanceGroupId": self.emr_task_instance_group_id,
"InstanceCount": 0
}])
def test_is_target_already_reached(self):
self.assertFalse(Emr.is_target_count_not_reached(0, 0))
def test_is_target_already_reached(self):
self.assertFalse(Emr.is_target_count_not_reached(0, 0))
