Add enable_availability_zone parameter to ec2 module.

cloud/ec2_elb

@@ -47,7 +47,7 @@ options:
     description:
       - AWS secret key. If not set then the value of the AWS_SECRET_KEY environment variable is used. 
     required: false
-    def2ault: None
+    default: None
     aliases: ['ec2_secret_key', 'secret_key' ]
   aws_access_key:
     description:
@@ -60,6 +60,13 @@ options:
       - The AWS region to use. If not specified then the value of the EC2_REGION environment variable, if any, is used.
     required: false
     aliases: ['aws_region', 'ec2_region']
+  enable_availability_zone:
+    description:
+      - Whether to enable the availability zone of the instance on the target ELB if the availability zone has not already
+        been enabled. If set to no, the task will fail if the availability zone is not enabled on the ELB.
+    required: false
+    default: yes
+    choices: [ "yes", "no" ]
   wait:
     description:
       - Wait for instance registration or deregistration to complete successfully before returning.  
@@ -106,6 +113,7 @@ AWS_REGIONS = ['ap-northeast-1',
 
 try:
     import boto
+    import boto.ec2
     import boto.ec2.elb
     from boto.regioninfo import RegionInfo
 except ImportError:
@@ -123,31 +131,36 @@ class ElbManager:
         self.instance_id = instance_id
         self.region = region
         self.lbs = self._get_instance_lbs(ec2_elbs)
-
+        self.changed = False
-        # if there are no ELBs to operate on
-        # there will be no changes made
-        if len(self.lbs) > 0:
-            self.changed = True
-        else:
-            self.changed = False
 
     def deregister(self, wait):
         """De-register the instance from all ELBs and wait for the ELB
         to report it out-of-service"""
 
         for lb in self.lbs:
+            initial_state = lb.get_instance_health([self.instance_id])[0]
             lb.deregister_instances([self.instance_id])
-            if wait: 
+            if wait:
-                self._await_elb_instance_state(lb, 'OutOfService')
+                self._await_elb_instance_state(lb, 'OutOfService', initial_state)
+            else:
+                # We cannot assume no change was made if we don't wait
+                # to find out
+                self.changed = True
 
-    def register(self, wait):
+    def register(self, wait, enable_availability_zone):
         """Register the instance for all ELBs and wait for the ELB
         to report the instance in-service"""
-
         for lb in self.lbs:
+            initial_state = lb.get_instance_health([self.instance_id])[0]
+            if enable_availability_zone:
+                self._enable_availailability_zone(lb)
             lb.register_instances([self.instance_id])
             if wait:
-                self._await_elb_instance_state(lb, 'InService')
+                self._await_elb_instance_state(lb, 'InService', initial_state)
+            else:
+                # We cannot assume no change was made if we don't wait
+                # to find out
+                self.changed = True
 
     def exists(self, lbtest):
         """ Verify that the named ELB actually exists """
@@ -159,16 +172,44 @@ class ElbManager:
                 break
         return found
 
+    def _enable_availailability_zone(self, lb):
+        """Enable the current instance's availability zone in the provided lb.
+        Returns True if the zone was enabled or False if no change was made.
+        lb: load balancer"""
+        instance = self._get_instance()
+        if instance.placement in lb.availability_zones:
+            return False
 
-    def _await_elb_instance_state(self, lb, awaited_state):
+        lb.enable_zones(zones=instance.placement)
+
+        # If successful, the new zone will have been added to
+        # lb.availability_zones
+        return instance.placement in lb.availability_zones
+
+    def _await_elb_instance_state(self, lb, awaited_state, initial_state):
         """Wait for an ELB to change state
         lb: load balancer
         awaited_state : state to poll for (string)"""
-
         while True:
-            state = lb.get_instance_health([self.instance_id])[0].state
+            instance_state = lb.get_instance_health([self.instance_id])[0]
-            if state == awaited_state:
+            if instance_state.state == awaited_state:
+                # Check the current state agains the initial state, and only set
+                # changed if they are different.
+                if instance_state.state != initial_state.state:
+                    self.changed = True
                 break
+            elif (awaited_state == 'InService'
+                  and instance_state.reason_code == "Instance"):
+                # If the reason_code for the instance being out of service is
+                # "Instance" this indicates a failure state, e.g. the instance
+                # has failed a health check or the ELB does not have the
+                # instance's availabilty zone enabled. The exact reason why is
+                # described in InstantState.description.
+                msg = ("The instance %s could not be put in service on %s."
+                       " Reason: %s")
+                self.module.fail_json(msg=msg % (self.instance_id,
+                                                 lb,
+                                                 instance_state.description))
             else:
                 time.sleep(1)
 
@@ -197,6 +238,16 @@ class ElbManager:
                         lbs.append(lb)
         return lbs
 
+    def _get_instance(self):
+        """Returns a boto.ec2.InstanceObject for self.instance_id"""
+        try:
+            endpoint = "ec2.%s.amazonaws.com" % self.region
+            connect_region = RegionInfo(name=self.region, endpoint=endpoint)
+            ec2_conn = boto.ec2.EC2Connection(self.aws_access_key, self.aws_secret_key, region=connect_region)
+        except boto.exception.NoAuthHandlerFound, e:
+            self.module.fail_json(msg=str(e))
+        return ec2_conn.get_only_instances(instance_ids=[self.instance_id])[0]
+
 
 def main():
 
@@ -209,7 +260,8 @@ def main():
             aws_secret_key={'default': None, 'aliases': ['ec2_secret_key', 'secret_key'], 'no_log': True},
             aws_access_key={'default': None, 'aliases': ['ec2_access_key', 'access_key']},
             region={'default': None, 'required': False, 'aliases':['aws_region', 'ec2_region'], 'choices':AWS_REGIONS},
-            wait={'required': False, 'choices': BOOLEANS, 'default': True}
+            enable_availability_zone={'default': True, 'required': False, 'choices': BOOLEANS, 'type': 'bool'},
+            wait={'required': False, 'choices': BOOLEANS, 'default': True, 'type': 'bool'}
         )
     )
 
@@ -218,6 +270,7 @@ def main():
     ec2_elbs = module.params['ec2_elbs']
     region = module.params['region']
     wait = module.params['wait']
+    enable_availability_zone = module.params['enable_availability_zone']
 
     if module.params['state'] == 'present' and 'ec2_elbs' not in module.params:
         module.fail_json(msg="ELBs are required for registration")
@@ -253,7 +306,7 @@ def main():
             module.fail_json(msg=msg)
 
     if module.params['state'] == 'present':
-        elb_man.register(wait)
+        elb_man.register(wait, enable_availability_zone)
     elif module.params['state'] == 'absent':
         elb_man.deregister(wait)