Merge pull request #202 from SHLo/support-custom-tags

Support for custom tagging #109
newsapps · Dec 20, 2017 · 4b8783d · 4b8783d
2 parents 85fd986 + beecddb
commit 4b8783d
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/beeswithmachineguns/bees.py b/beeswithmachineguns/bees.py
@@ -39,6 +39,7 @@
 import socket
 import time
 import sys
+import ast
 IS_PY2 = sys.version_info.major == 2
 if IS_PY2:
     from urllib.request import urlopen, Request
@@ -128,7 +129,7 @@ def _get_security_group_id(connection, security_group_name, subnet):
 
 # Methods
 
-def up(count, group, zone, image_id, instance_type, username, key_name, subnet, bid = None):
+def up(count, group, zone, image_id, instance_type, username, key_name, subnet, tags, bid = None):
     """
     Startup the load testing server.
     """
@@ -224,6 +225,14 @@ def up(count, group, zone, image_id, instance_type, username, key_name, subnet,
             return e
 
         instances = reservation.instances
+    if tags:
+        try:
+            tags_dict = ast.literal_eval(tags)
+            ids = [instance.id for instance in instances]
+            ec2_connection.create_tags(ids, tags_dict)
+        except Exception as e:
+            print("Unable to create tags:")
+            print("example: bees up -x \"{'any_key': 'any_value'}\"")
 
     if instance_ids:
         existing_reservations = ec2_connection.get_all_instances(instance_ids=instance_ids)

diff --git a/beeswithmachineguns/main.py b/beeswithmachineguns/main.py
@@ -89,6 +89,9 @@ def parse_options():
     up_group.add_option('-b', '--bid', metavar="BID", nargs=1,
                         action='store', dest='bid', type='float', default=None,
                         help="The maximum bid price per spot instance (default: None).")
+    up_group.add_option('-x', '--tags', metavar="TAGS", nargs=1,
+                        action='store', dest='tags', type='string', default=None,
+                        help="custome tags for bee instances")
 
     parser.add_option_group(up_group)
 
@@ -206,11 +209,11 @@ def parse_options():
                                                             options.zone, options.instance,
                                                             options.type,options.login,
                                                             options.key, options.subnet,
-                                                            options.bid)).start()
+                                                            options.tags, options.bid)).start()
                     #time allowed between threads
                     time.sleep(delay)
         else:
-            bees.up(options.servers, options.group, options.zone, options.instance, options.type, options.login, options.key, options.subnet, options.bid)
+            bees.up(options.servers, options.group, options.zone, options.instance, options.type, options.login, options.key, options.subnet, options.tags, options.bid)
 
     elif command == 'attack':
         if not options.url: