From e296c7df9d9edcd727a095b3d11677e3a491d8c6 Mon Sep 17 00:00:00 2001
From: Adam Keller <akeller@dandb.com>
Date: Wed, 14 Oct 2015 13:32:53 -0700
Subject: [PATCH] Added ami_tag method

---
 License2Deploy/rolling_deploy.py | 17 ++++++++++++++++-
 tests/rolling_deploy_test.py     |  9 +++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/License2Deploy/rolling_deploy.py b/License2Deploy/rolling_deploy.py
index 4f03507..6f31011 100644
--- a/License2Deploy/rolling_deploy.py
+++ b/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
diff --git a/tests/rolling_deploy_test.py b/tests/rolling_deploy_test.py
index 994efcb..5466472 100644
--- a/tests/rolling_deploy_test.py
+++ b/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')