Added ami_tag method

License2Deploy/rolling_deploy.py

@@ -166,7 +166,7 @@ def lb_healthcheck(self, new_ids, retry=10, wait_time=10):
       logging.info("ELB healthcheck OK == {0}: {1}".format(instance_id.instance_id, instance_id.state))
     return True
 
-  def confirm_lb_has_only_new_instances(self, wait_time=30):
+  def confirm_lb_has_only_new_instances(self, wait_time=60):
     ''' Confirm that only new instances with the current build tag are in the load balancer '''
     sleep(wait_time) # Allotting time for the instances to shut down
     lb = self.get_lb()
@@ -179,6 +179,20 @@ def confirm_lb_has_only_new_instances(self, wait_time=30):
     logging.info("Deployed instances {0} to ELB: {1}".format(instance_ids, lb))
     return instance_ids
 
+  def tag_ami(self, ami_id, env):
+    ''' Tagging AMI with DEPLOYED tag '''
+    try:
+      current_tag = self.conn_ec2.get_all_images(image_ids=ami_id)[0].tags.get('DEPLOYED')
+      if not current_tag: 
+        self.conn_ec2.create_tags([self.ami_id], {"deployed": env})
+        logging.info("No DEPLOY tags exist, tagged with {0}".format(env))
+      elif env not in current_tag:
+        new_tag = ', '.join([current_tag, env])
+        self.conn_ec2.create_tags([self.ami_id], {"deployed": new_tag})
+        logging.info("DEPLOY tags currently exist: {0}, new tag is {0}".format(current_tag, new_tag))
+    except Exception as e:
+      logging.error("Unable to tag ID, please investigate: {0}".format(e))
+
   def healthcheck_new_instances(self, group_name): # pragma: no cover
     ''' Healthchecking new instances to ensure deployment was successful '''
     instance_ids = self.get_all_instance_ids(group_name)
@@ -197,6 +211,7 @@ def deploy(self): # pragma: no cover
     self.healthcheck_new_instances(group_name)
     self.set_autoscale_instance_desired_count(self.calculate_autoscale_desired_instance_count(group_name, 'decrease'), group_name)
     self.confirm_lb_has_only_new_instances()
+    self.tag_ami(self.ami_id, self.env)
     logging.info("Deployment Complete!")
 
 def get_args(): # pragma: no cover

tests/rolling_deploy_test.py

@@ -74,6 +74,15 @@ def setUpEC2(self):
 
     return [conn, instance_id_list]
 
+  @mock_ec2
+  def test_tag_ami(self):
+    conn = self.setUpEC2()[0]
+    inst_ids = self.setUpEC2()[1]
+    conn.create_image(inst_ids[0], "test-ami", "this is a test ami")
+    ami_ids = conn.get_all_images()
+    ami_id = ami_ids[0].id
+    self.assertEqual(self.rolling_deploy.tag_ami(ami_id, 'qa'), None)
+
   @mock_ec2
   def test_load_config(self):
     self.assertEqual(AWSConn.load_config('regions.yml').get('qa'), 'us-west-1')