Merge pull request #967 from JackDanger/elb-timeout-is-not-required

ELB connection draining timeout defaults to 300 seconds
getmoto · May 31, 2017 · 8f7f306 · 8f7f306
2 parents d8a8f88 + b0c83c4
commit 8f7f306
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 7 deletions.
diff --git a/moto/elb/responses.py b/moto/elb/responses.py
@@ -159,9 +159,8 @@ def modify_load_balancer_attributes(self):
         if connection_draining:
             attribute = ConnectionDrainingAttribute()
             attribute.enabled = connection_draining["enabled"] == "true"
-            attribute.timeout = connection_draining["timeout"]
-            self.elb_backend.set_connection_draining_attribute(
-                load_balancer_name, attribute)
+            attribute.timeout = connection_draining.get("timeout", 300)
+            self.elb_backend.set_connection_draining_attribute(load_balancer_name, attribute)
 
         connection_settings = self._get_dict_param(
             "LoadBalancerAttributes.ConnectionSettings.")
@@ -172,7 +171,7 @@ def modify_load_balancer_attributes(self):
                 load_balancer_name, attribute)
 
         template = self.response_template(MODIFY_ATTRIBUTES_TEMPLATE)
-        return template.render(attributes=load_balancer.attributes)
+        return template.render(load_balancer=load_balancer, attributes=load_balancer.attributes)
 
     def create_load_balancer_policy(self):
         load_balancer_name = self._get_param('LoadBalancerName')
@@ -592,9 +591,11 @@ def _add_tags(self, elb):
         <Enabled>{{ attributes.cross_zone_load_balancing.enabled }}</Enabled>
       </CrossZoneLoadBalancing>
       <ConnectionDraining>
-        <Enabled>{{ attributes.connection_draining.enabled }}</Enabled>
         {% if attributes.connection_draining.enabled %}
+        <Enabled>true</Enabled>
         <Timeout>{{ attributes.connection_draining.timeout }}</Timeout>
+        {% else %}
+        <Enabled>false</Enabled>
         {% endif %}
       </ConnectionDraining>
     </LoadBalancerAttributes>
@@ -607,7 +608,7 @@ def _add_tags(self, elb):
 
 MODIFY_ATTRIBUTES_TEMPLATE = """<ModifyLoadBalancerAttributesResponse xmlns="http://elasticloadbalancing.amazonaws.com/doc/2012-06-01/">
   <ModifyLoadBalancerAttributesResult>
-  <LoadBalancerName>my-loadbalancer</LoadBalancerName>
+  <LoadBalancerName>{{ load_balancer.name }}</LoadBalancerName>
     <LoadBalancerAttributes>
       <AccessLog>
         <Enabled>{{ attributes.access_log.enabled }}</Enabled>
@@ -624,9 +625,11 @@ def _add_tags(self, elb):
         <Enabled>{{ attributes.cross_zone_load_balancing.enabled }}</Enabled>
       </CrossZoneLoadBalancing>
       <ConnectionDraining>
-        <Enabled>{{ attributes.connection_draining.enabled }}</Enabled>
         {% if attributes.connection_draining.enabled %}
+        <Enabled>true</Enabled>
         <Timeout>{{ attributes.connection_draining.timeout }}</Timeout>
+        {% else %}
+        <Enabled>false</Enabled>
         {% endif %}
       </ConnectionDraining>
     </LoadBalancerAttributes>

diff --git a/tests/test_elb/test_elb.py b/tests/test_elb/test_elb.py
@@ -816,6 +816,42 @@ def test_create_with_tags():
     tags.should.have.key('k').which.should.equal('v')
 
 
+@mock_elb
+def test_modify_attributes():
+    client = boto3.client('elb', region_name='us-east-1')
+
+    client.create_load_balancer(
+        LoadBalancerName='my-lb',
+        Listeners=[{'Protocol': 'tcp', 'LoadBalancerPort': 80, 'InstancePort': 8080}],
+        AvailabilityZones=['us-east-1a', 'us-east-1b']
+    )
+
+    # Default ConnectionDraining timeout of 300 seconds
+    client.modify_load_balancer_attributes(
+        LoadBalancerName='my-lb',
+        LoadBalancerAttributes={
+            'ConnectionDraining': {'Enabled': True},
+        }
+    )
+    lb_attrs = client.describe_load_balancer_attributes(LoadBalancerName='my-lb')
+    lb_attrs['LoadBalancerAttributes']['ConnectionDraining']['Enabled'].should.equal(True)
+    lb_attrs['LoadBalancerAttributes']['ConnectionDraining']['Timeout'].should.equal(300)
+
+    # specify a custom ConnectionDraining timeout
+    client.modify_load_balancer_attributes(
+        LoadBalancerName='my-lb',
+        LoadBalancerAttributes={
+            'ConnectionDraining': {
+                'Enabled': True,
+                'Timeout': 45,
+            },
+        }
+    )
+    lb_attrs = client.describe_load_balancer_attributes(LoadBalancerName='my-lb')
+    lb_attrs['LoadBalancerAttributes']['ConnectionDraining']['Enabled'].should.equal(True)
+    lb_attrs['LoadBalancerAttributes']['ConnectionDraining']['Timeout'].should.equal(45)
+
+
 @mock_ec2
 @mock_elb
 def test_subnets():