Merge pull request #1284 from carleton-stripe/asg-health-status

Add support for HealthStatus for Auto Scaling Groups

moto/autoscaling/models.py

@@ -16,9 +16,10 @@ ASG_NAME_TAG = "aws:autoscaling:groupName"
 
 
 class InstanceState(object):
-    def __init__(self, instance, lifecycle_state="InService"):
+    def __init__(self, instance, lifecycle_state="InService", health_status="Healthy"):
         self.instance = instance
         self.lifecycle_state = lifecycle_state
+        self.health_status = health_status
 
 
 class FakeScalingPolicy(BaseModel):
@@ -432,6 +433,12 @@ class AutoScalingBackend(BaseBackend):
             group.instance_states.extend(new_instances)
             self.update_attached_elbs(group.name)
 
+    def set_instance_health(self, instance_id, health_status, should_respect_grace_period):
+        instance = self.ec2_backend.get_instance(instance_id)
+        instance_state = next(instance_state for group in self.autoscaling_groups.values()
+                              for instance_state in group.instance_states if instance_state.instance.id == instance.id)
+        instance_state.health_status = health_status
+
     def detach_instances(self, group_name, instance_ids, should_decrement):
         group = self.autoscaling_groups[group_name]
         original_size = len(group.instance_states)

moto/autoscaling/responses.py

@@ -98,6 +98,18 @@ class AutoScalingResponse(BaseResponse):
         template = self.response_template(ATTACH_INSTANCES_TEMPLATE)
         return template.render()
 
+    @amz_crc32
+    @amzn_request_id
+    def set_instance_health(self):
+        instance_id = self._get_param('InstanceId')
+        health_status = self._get_param("HealthStatus")
+        if health_status not in ['Healthy', 'Unhealthy']:
+            raise ValueError('Valid instance health states are: [Healthy, Unhealthy]')
+        should_respect_grace_period = self._get_param("ShouldRespectGracePeriod")
+        self.autoscaling_backend.set_instance_health(instance_id, health_status, should_respect_grace_period)
+        template = self.response_template(SET_INSTANCE_HEALTH_TEMPLATE)
+        return template.render()
+
     @amz_crc32
     @amzn_request_id
     def detach_instances(self):
@@ -397,7 +409,7 @@ DESCRIBE_AUTOSCALING_GROUPS_TEMPLATE = """<DescribeAutoScalingGroupsResponse xml
         <Instances>
           {% for instance_state in group.instance_states %}
           <member>
-            <HealthStatus>HEALTHY</HealthStatus>
+            <HealthStatus>{{ instance_state.health_status }}</HealthStatus>
             <AvailabilityZone>us-east-1e</AvailabilityZone>
             <InstanceId>{{ instance_state.instance.id }}</InstanceId>
             <LaunchConfigurationName>{{ group.launch_config_name }}</LaunchConfigurationName>
@@ -472,7 +484,7 @@ DESCRIBE_AUTOSCALING_INSTANCES_TEMPLATE = """<DescribeAutoScalingInstancesRespon
     <AutoScalingInstances>
       {% for instance_state in instance_states %}
       <member>
-        <HealthStatus>HEALTHY</HealthStatus>
+        <HealthStatus>{{ instance_state.health_status }}</HealthStatus>
         <AutoScalingGroupName>{{ instance_state.instance.autoscaling_group.name }}</AutoScalingGroupName>
         <AvailabilityZone>us-east-1e</AvailabilityZone>
         <InstanceId>{{ instance_state.instance.id }}</InstanceId>
@@ -568,3 +580,10 @@ DETACH_LOAD_BALANCERS_TEMPLATE = """<DetachLoadBalancersResponse xmlns="http://a
 <RequestId>{{ requestid }}</RequestId>
 </ResponseMetadata>
 </DetachLoadBalancersResponse>"""
+
+SET_INSTANCE_HEALTH_TEMPLATE = """<SetInstanceHealthResponse xmlns="http://autoscaling.amazonaws.com/doc/2011-01-01/">
+<SetInstanceHealthResponse></SetInstanceHealthResponse>
+<ResponseMetadata>
+<RequestId>{{ requestid }}</RequestId>
+</ResponseMetadata>
+</SetInstanceHealthResponse>"""

tests/test_autoscaling/test_autoscaling.py

@@ -311,6 +311,7 @@ def test_autoscaling_group_describe_instances():
     instances = list(conn.get_all_autoscaling_instances())
     instances.should.have.length_of(2)
     instances[0].launch_config_name.should.equal('tester')
+    instances[0].health_status.should.equal('Healthy')
     autoscale_instance_ids = [instance.instance_id for instance in instances]
 
     ec2_conn = boto.connect_ec2()
@@ -959,3 +960,56 @@ def test_attach_one_instance():
         AutoScalingGroupNames=['test_asg']
     )
     response['AutoScalingGroups'][0]['Instances'].should.have.length_of(3)
+
+@mock_autoscaling
+@mock_ec2
+def test_describe_instance_health():
+    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=2,
+        MaxSize=4,
+        DesiredCapacity=2,
+    )
+
+    response = client.describe_auto_scaling_groups(
+        AutoScalingGroupNames=['test_asg']
+    )
+
+    instance1 = response['AutoScalingGroups'][0]['Instances'][0]
+    instance1['HealthStatus'].should.equal('Healthy')
+
+@mock_autoscaling
+@mock_ec2
+def test_set_instance_health():
+    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=2,
+        MaxSize=4,
+        DesiredCapacity=2,
+    )
+
+    response = client.describe_auto_scaling_groups(
+        AutoScalingGroupNames=['test_asg']
+    )
+
+    instance1 = response['AutoScalingGroups'][0]['Instances'][0]
+    instance1['HealthStatus'].should.equal('Healthy')
+
+    client.set_instance_health(InstanceId=instance1['InstanceId'], HealthStatus='Unhealthy')
+
+    response = client.describe_auto_scaling_groups(
+        AutoScalingGroupNames=['test_asg']
+    )
+
+    instance1 = response['AutoScalingGroups'][0]['Instances'][0]
+    instance1['HealthStatus'].should.equal('Unhealthy')