Add instance protection support in autoscaling, with tests

This commit is contained in:
Ka Wai Wan 2018-11-24 02:32:39 -08:00
parent ed861ecae1
commit df2120f38c
3 changed files with 228 additions and 22 deletions

View File

@ -17,10 +17,12 @@ ASG_NAME_TAG = "aws:autoscaling:groupName"
class InstanceState(object):
def __init__(self, instance, lifecycle_state="InService", health_status="Healthy"):
def __init__(self, instance, lifecycle_state="InService",
health_status="Healthy", protected_from_scale_in=False):
self.instance = instance
self.lifecycle_state = lifecycle_state
self.health_status = health_status
self.protected_from_scale_in = protected_from_scale_in
class FakeScalingPolicy(BaseModel):
@ -152,7 +154,8 @@ class FakeAutoScalingGroup(BaseModel):
min_size, launch_config_name, vpc_zone_identifier,
default_cooldown, health_check_period, health_check_type,
load_balancers, target_group_arns, placement_group, termination_policies,
autoscaling_backend, tags):
autoscaling_backend, tags,
new_instances_protected_from_scale_in=False):
self.autoscaling_backend = autoscaling_backend
self.name = name
@ -178,6 +181,7 @@ class FakeAutoScalingGroup(BaseModel):
self.target_group_arns = target_group_arns
self.placement_group = placement_group
self.termination_policies = termination_policies
self.new_instances_protected_from_scale_in = new_instances_protected_from_scale_in
self.suspended_processes = []
self.instance_states = []
@ -210,6 +214,8 @@ class FakeAutoScalingGroup(BaseModel):
placement_group=None,
termination_policies=properties.get("TerminationPolicies", []),
tags=properties.get("Tags", []),
new_instances_protected_from_scale_in=properties.get(
"NewInstancesProtectedFromScaleIn", False)
)
return group
@ -238,7 +244,8 @@ class FakeAutoScalingGroup(BaseModel):
def update(self, availability_zones, desired_capacity, max_size, min_size,
launch_config_name, vpc_zone_identifier, default_cooldown,
health_check_period, health_check_type,
placement_group, termination_policies):
placement_group, termination_policies,
new_instances_protected_from_scale_in=None):
if availability_zones:
self.availability_zones = availability_zones
if max_size is not None:
@ -256,6 +263,8 @@ class FakeAutoScalingGroup(BaseModel):
self.health_check_period = health_check_period
if health_check_type is not None:
self.health_check_type = health_check_type
if new_instances_protected_from_scale_in is not None:
self.new_instances_protected_from_scale_in = new_instances_protected_from_scale_in
if desired_capacity is not None:
self.set_desired_capacity(desired_capacity)
@ -280,12 +289,16 @@ class FakeAutoScalingGroup(BaseModel):
else:
# Need to remove some instances
count_to_remove = curr_instance_count - self.desired_capacity
instances_to_remove = self.instance_states[:count_to_remove]
instance_ids_to_remove = [
instance.instance.id for instance in instances_to_remove]
self.autoscaling_backend.ec2_backend.terminate_instances(
instance_ids_to_remove)
self.instance_states = self.instance_states[count_to_remove:]
instances_to_remove = [ # only remove unprotected
state for state in self.instance_states
if not state.protected_from_scale_in
][:count_to_remove]
if instances_to_remove: # just in case not instances to remove
instance_ids_to_remove = [
instance.instance.id for instance in instances_to_remove]
self.autoscaling_backend.ec2_backend.terminate_instances(
instance_ids_to_remove)
self.instance_states = list(set(self.instance_states) - set(instances_to_remove))
def get_propagated_tags(self):
propagated_tags = {}
@ -310,7 +323,10 @@ class FakeAutoScalingGroup(BaseModel):
)
for instance in reservation.instances:
instance.autoscaling_group = self
self.instance_states.append(InstanceState(instance))
self.instance_states.append(InstanceState(
instance,
protected_from_scale_in=self.new_instances_protected_from_scale_in,
))
def append_target_groups(self, target_group_arns):
append = [x for x in target_group_arns if x not in self.target_group_arns]
@ -372,7 +388,8 @@ class AutoScalingBackend(BaseBackend):
default_cooldown, health_check_period,
health_check_type, load_balancers,
target_group_arns, placement_group,
termination_policies, tags):
termination_policies, tags,
new_instances_protected_from_scale_in=False):
def make_int(value):
return int(value) if value is not None else value
@ -403,6 +420,7 @@ class AutoScalingBackend(BaseBackend):
termination_policies=termination_policies,
autoscaling_backend=self,
tags=tags,
new_instances_protected_from_scale_in=new_instances_protected_from_scale_in,
)
self.autoscaling_groups[name] = group
@ -415,12 +433,14 @@ class AutoScalingBackend(BaseBackend):
launch_config_name, vpc_zone_identifier,
default_cooldown, health_check_period,
health_check_type, placement_group,
termination_policies):
termination_policies,
new_instances_protected_from_scale_in=None):
group = self.autoscaling_groups[name]
group.update(availability_zones, desired_capacity, max_size,
min_size, launch_config_name, vpc_zone_identifier,
default_cooldown, health_check_period, health_check_type,
placement_group, termination_policies)
placement_group, termination_policies,
new_instances_protected_from_scale_in=new_instances_protected_from_scale_in)
return group
def describe_auto_scaling_groups(self, names):
@ -448,7 +468,13 @@ class AutoScalingBackend(BaseBackend):
raise ResourceContentionError
else:
group.desired_capacity = original_size + len(instance_ids)
new_instances = [InstanceState(self.ec2_backend.get_instance(x)) for x in instance_ids]
new_instances = [
InstanceState(
self.ec2_backend.get_instance(x),
protected_from_scale_in=group.new_instances_protected_from_scale_in,
)
for x in instance_ids
]
for instance in new_instances:
self.ec2_backend.create_tags([instance.instance.id], {ASG_NAME_TAG: group.name})
group.instance_states.extend(new_instances)
@ -626,6 +652,13 @@ class AutoScalingBackend(BaseBackend):
group = self.autoscaling_groups[group_name]
group.suspended_processes = scaling_processes or []
def set_instance_protection(self, group_name, instance_ids, protected_from_scale_in):
group = self.autoscaling_groups[group_name]
protected_instances = [
x for x in group.instance_states if x.instance.id in instance_ids]
for instance in protected_instances:
instance.protected_from_scale_in = protected_from_scale_in
autoscaling_backends = {}
for region, ec2_backend in ec2_backends.items():

View File

@ -85,6 +85,8 @@ class AutoScalingResponse(BaseResponse):
termination_policies=self._get_multi_param(
'TerminationPolicies.member'),
tags=self._get_list_prefix('Tags.member'),
new_instances_protected_from_scale_in=self._get_bool_param(
'NewInstancesProtectedFromScaleIn', False)
)
template = self.response_template(CREATE_AUTOSCALING_GROUP_TEMPLATE)
return template.render()
@ -192,6 +194,8 @@ class AutoScalingResponse(BaseResponse):
placement_group=self._get_param('PlacementGroup'),
termination_policies=self._get_multi_param(
'TerminationPolicies.member'),
new_instances_protected_from_scale_in=self._get_bool_param(
'NewInstancesProtectedFromScaleIn', None)
)
template = self.response_template(UPDATE_AUTOSCALING_GROUP_TEMPLATE)
return template.render()
@ -290,6 +294,15 @@ class AutoScalingResponse(BaseResponse):
template = self.response_template(SUSPEND_PROCESSES_TEMPLATE)
return template.render()
def set_instance_protection(self):
group_name = self._get_param('AutoScalingGroupName')
instance_ids = self._get_multi_param('InstanceIds.member')
protected_from_scale_in = self._get_bool_param('ProtectedFromScaleIn')
self.autoscaling_backend.set_instance_protection(
group_name, instance_ids, protected_from_scale_in)
template = self.response_template(SET_INSTANCE_PROTECTION_TEMPLATE)
return template.render()
CREATE_LAUNCH_CONFIGURATION_TEMPLATE = """<CreateLaunchConfigurationResponse xmlns="http://autoscaling.amazonaws.com/doc/2011-01-01/">
<ResponseMetadata>
@ -490,6 +503,7 @@ DESCRIBE_AUTOSCALING_GROUPS_TEMPLATE = """<DescribeAutoScalingGroupsResponse xml
<InstanceId>{{ instance_state.instance.id }}</InstanceId>
<LaunchConfigurationName>{{ group.launch_config_name }}</LaunchConfigurationName>
<LifecycleState>{{ instance_state.lifecycle_state }}</LifecycleState>
<ProtectedFromScaleIn>{{ instance_state.protected_from_scale_in|string|lower }}</ProtectedFromScaleIn>
</member>
{% endfor %}
</Instances>
@ -530,6 +544,7 @@ DESCRIBE_AUTOSCALING_GROUPS_TEMPLATE = """<DescribeAutoScalingGroupsResponse xml
{% if group.placement_group %}
<PlacementGroup>{{ group.placement_group }}</PlacementGroup>
{% endif %}
<NewInstancesProtectedFromScaleIn>{{ group.new_instances_protected_from_scale_in|string|lower }}</NewInstancesProtectedFromScaleIn>
</member>
{% endfor %}
</AutoScalingGroups>
@ -565,6 +580,7 @@ DESCRIBE_AUTOSCALING_INSTANCES_TEMPLATE = """<DescribeAutoScalingInstancesRespon
<InstanceId>{{ instance_state.instance.id }}</InstanceId>
<LaunchConfigurationName>{{ instance_state.instance.autoscaling_group.launch_config_name }}</LaunchConfigurationName>
<LifecycleState>{{ instance_state.lifecycle_state }}</LifecycleState>
<ProtectedFromScaleIn>{{ instance_state.protected_from_scale_in|string|lower }}</ProtectedFromScaleIn>
</member>
{% endfor %}
</AutoScalingInstances>
@ -668,3 +684,10 @@ SET_INSTANCE_HEALTH_TEMPLATE = """<SetInstanceHealthResponse xmlns="http://autos
<RequestId>{{ requestid }}</RequestId>
</ResponseMetadata>
</SetInstanceHealthResponse>"""
SET_INSTANCE_PROTECTION_TEMPLATE = """<SetInstanceProtectionResponse xmlns="http://autoscaling.amazonaws.com/doc/2011-01-01/">
<SetInstanceProtectionResult></SetInstanceProtectionResult>
<ResponseMetadata>
<RequestId>{{ requestid }}</RequestId>
</ResponseMetadata>
</SetInstanceProtectionResponse>"""

View File

@ -710,6 +710,7 @@ def test_create_autoscaling_group_boto3():
'PropagateAtLaunch': False
}],
VPCZoneIdentifier=mocked_networking['subnet1'],
NewInstancesProtectedFromScaleIn=False,
)
response['ResponseMetadata']['HTTPStatusCode'].should.equal(200)
@ -728,13 +729,48 @@ def test_describe_autoscaling_groups_boto3():
MaxSize=20,
DesiredCapacity=5,
VPCZoneIdentifier=mocked_networking['subnet1'],
NewInstancesProtectedFromScaleIn=True,
)
response = client.describe_auto_scaling_groups(
AutoScalingGroupNames=["test_asg"]
)
response['ResponseMetadata']['HTTPStatusCode'].should.equal(200)
response['AutoScalingGroups'][0][
'AutoScalingGroupName'].should.equal('test_asg')
group = response['AutoScalingGroups'][0]
group['AutoScalingGroupName'].should.equal('test_asg')
group['NewInstancesProtectedFromScaleIn'].should.equal(True)
group['Instances'][0]['ProtectedFromScaleIn'].should.equal(True)
@mock_autoscaling
def test_describe_autoscaling_instances_boto3():
mocked_networking = setup_networking()
client = boto3.client('autoscaling', region_name='us-east-1')
_ = client.create_launch_configuration(
LaunchConfigurationName='test_launch_configuration'
)
_ = client.create_auto_scaling_group(
AutoScalingGroupName='test_asg',
LaunchConfigurationName='test_launch_configuration',
MinSize=0,
MaxSize=20,
DesiredCapacity=5,
VPCZoneIdentifier=mocked_networking['subnet1'],
NewInstancesProtectedFromScaleIn=True,
)
response = client.describe_auto_scaling_groups(
AutoScalingGroupNames=["test_asg"]
)
instance_ids = [
instance['InstanceId']
for instance in response['AutoScalingGroups'][0]['Instances']
]
response = client.describe_auto_scaling_instances(InstanceIds=instance_ids)
for instance in response['AutoScalingInstances']:
instance['AutoScalingGroupName'].should.equal('test_asg')
instance['ProtectedFromScaleIn'].should.equal(True)
@mock_autoscaling
@ -751,17 +787,21 @@ def test_update_autoscaling_group_boto3():
MaxSize=20,
DesiredCapacity=5,
VPCZoneIdentifier=mocked_networking['subnet1'],
NewInstancesProtectedFromScaleIn=True,
)
response = client.update_auto_scaling_group(
_ = client.update_auto_scaling_group(
AutoScalingGroupName='test_asg',
MinSize=1,
NewInstancesProtectedFromScaleIn=False,
)
response = client.describe_auto_scaling_groups(
AutoScalingGroupNames=["test_asg"]
)
response['AutoScalingGroups'][0]['MinSize'].should.equal(1)
group = response['AutoScalingGroups'][0]
group['MinSize'].should.equal(1)
group['NewInstancesProtectedFromScaleIn'].should.equal(False)
@mock_autoscaling
@ -992,9 +1032,7 @@ def test_attach_one_instance():
'PropagateAtLaunch': True
}],
VPCZoneIdentifier=mocked_networking['subnet1'],
)
response = client.describe_auto_scaling_groups(
AutoScalingGroupNames=['test_asg']
NewInstancesProtectedFromScaleIn=True,
)
ec2 = boto3.resource('ec2', 'us-east-1')
@ -1009,7 +1047,11 @@ def test_attach_one_instance():
response = client.describe_auto_scaling_groups(
AutoScalingGroupNames=['test_asg']
)
response['AutoScalingGroups'][0]['Instances'].should.have.length_of(3)
instances = response['AutoScalingGroups'][0]['Instances']
instances.should.have.length_of(3)
for instance in instances:
instance['ProtectedFromScaleIn'].should.equal(True)
@mock_autoscaling
@mock_ec2
@ -1100,3 +1142,111 @@ def test_suspend_processes():
launch_suspended = True
assert launch_suspended is True
@mock_autoscaling
def test_set_instance_protection():
mocked_networking = setup_networking()
client = boto3.client('autoscaling', region_name='us-east-1')
_ = client.create_launch_configuration(
LaunchConfigurationName='test_launch_configuration'
)
_ = client.create_auto_scaling_group(
AutoScalingGroupName='test_asg',
LaunchConfigurationName='test_launch_configuration',
MinSize=0,
MaxSize=20,
DesiredCapacity=5,
VPCZoneIdentifier=mocked_networking['subnet1'],
NewInstancesProtectedFromScaleIn=False,
)
response = client.describe_auto_scaling_groups(AutoScalingGroupNames=['test_asg'])
instance_ids = [
instance['InstanceId']
for instance in response['AutoScalingGroups'][0]['Instances']
]
protected = instance_ids[:3]
_ = client.set_instance_protection(
AutoScalingGroupName='test_asg',
InstanceIds=protected,
ProtectedFromScaleIn=True,
)
response = client.describe_auto_scaling_groups(AutoScalingGroupNames=['test_asg'])
for instance in response['AutoScalingGroups'][0]['Instances']:
instance['ProtectedFromScaleIn'].should.equal(
instance['InstanceId'] in protected
)
@mock_autoscaling
def test_set_desired_capacity_up_boto3():
mocked_networking = setup_networking()
client = boto3.client('autoscaling', region_name='us-east-1')
_ = client.create_launch_configuration(
LaunchConfigurationName='test_launch_configuration'
)
_ = client.create_auto_scaling_group(
AutoScalingGroupName='test_asg',
LaunchConfigurationName='test_launch_configuration',
MinSize=0,
MaxSize=20,
DesiredCapacity=5,
VPCZoneIdentifier=mocked_networking['subnet1'],
NewInstancesProtectedFromScaleIn=True,
)
_ = client.set_desired_capacity(
AutoScalingGroupName='test_asg',
DesiredCapacity=10,
)
response = client.describe_auto_scaling_groups(AutoScalingGroupNames=['test_asg'])
instances = response['AutoScalingGroups'][0]['Instances']
instances.should.have.length_of(10)
for instance in instances:
instance['ProtectedFromScaleIn'].should.equal(True)
@mock_autoscaling
def test_set_desired_capacity_down_boto3():
mocked_networking = setup_networking()
client = boto3.client('autoscaling', region_name='us-east-1')
_ = client.create_launch_configuration(
LaunchConfigurationName='test_launch_configuration'
)
_ = client.create_auto_scaling_group(
AutoScalingGroupName='test_asg',
LaunchConfigurationName='test_launch_configuration',
MinSize=0,
MaxSize=20,
DesiredCapacity=5,
VPCZoneIdentifier=mocked_networking['subnet1'],
NewInstancesProtectedFromScaleIn=True,
)
response = client.describe_auto_scaling_groups(AutoScalingGroupNames=['test_asg'])
instance_ids = [
instance['InstanceId']
for instance in response['AutoScalingGroups'][0]['Instances']
]
unprotected, protected = instance_ids[:2], instance_ids[2:]
_ = client.set_instance_protection(
AutoScalingGroupName='test_asg',
InstanceIds=unprotected,
ProtectedFromScaleIn=False,
)
_ = client.set_desired_capacity(
AutoScalingGroupName='test_asg',
DesiredCapacity=1,
)
response = client.describe_auto_scaling_groups(AutoScalingGroupNames=['test_asg'])
group = response['AutoScalingGroups'][0]
group['DesiredCapacity'].should.equal(1)
instance_ids = {instance['InstanceId'] for instance in group['Instances']}
set(protected).should.equal(instance_ids)
set(unprotected).should_not.be.within(instance_ids) # only unprotected killed