From 041039359b9966bc12802ff83b4b31ef356decb6 Mon Sep 17 00:00:00 2001
From: Yichen Wang <yicwang@cisco.com>
Date: Mon, 7 Dec 2015 13:01:21 -0800
Subject: [PATCH] Add more features for Mercury

1. Add --delete-image-after-run;
2. Add --availability-zone to CLI options;
3. Doc updates;

Change-Id: If26d8c60e73ef119c70fad815cdf8b11cf23a001
---
 doc/source/usage.rst | 13 ++++++++++---
 vmtp/compute.py      | 11 +++++++++++
 vmtp/vmtp.py         | 24 ++++++++++++++++++++++--
 3 files changed, 43 insertions(+), 5 deletions(-)

diff --git a/doc/source/usage.rst b/doc/source/usage.rst
index f5cbb2d..1951e38 100644
--- a/doc/source/usage.rst
+++ b/doc/source/usage.rst
@@ -13,12 +13,14 @@ VMTP Usage
                    [--external-host <user>@<host_ssh_ip>[:password>]]
                    [--controller-node <user>@<host_ssh_ip>[:<password>]]
                    [--mongod-server <server ip>] [--json <file>]
-                   [--tp-tool <nuttcp|iperf>] [--hypervisor [<az>:] <hostname>]
-                   [--inter-node-only] [--protocols <T|U|I>]
+                   [--tp-tool <nuttcp|iperf>]
+                   [--availability_zone <availability_zone>] [--hypervisor [<az>:]
+                   <hostname>] [--inter-node-only] [--protocols <T|U|I>]
                    [--bandwidth <bandwidth>] [--tcpbuf <tcp_pkt_size1,...>]
                    [--udpbuf <udp_pkt_size1,...>]
                    [--reuse_network_name <network_name>]
-                   [--os-dataplane-network <network_name>] [--no-env]
+                   [--os-dataplane-network <network_name>]
+                   [--delete-image-after-run] [--no-env]
                    [--vnic-type <direct|macvtap|normal>] [-d] [-v]
                    [--stop-on-error] [--vm-image-url <url_to_image>]
                    [--test-description <test_description>]
@@ -49,6 +51,8 @@ VMTP Usage
       --json <file>         store results in json format file
       --tp-tool <nuttcp|iperf>
                             transport perf tool to use (default=nuttcp)
+      --availability_zone <availability_zone>
+                            availability zone for running VMTP
       --hypervisor [<az>:] <hostname>
                             hypervisor to use (1 per arg, up to 2 args)
       --inter-node-only     only measure inter-node
@@ -67,6 +71,9 @@ VMTP Usage
       --os-dataplane-network <network_name>
                             Internal network name for OpenStack to hold data plane
                             traffic
+      --delete-image-after-run
+                            delete image that are uploaded by VMTP when tests are
+                            finished
       --no-env              do not read env variables
       --vnic-type <direct|macvtap|normal>
                             binding vnic type for test VMs
diff --git a/vmtp/compute.py b/vmtp/compute.py
index ef8593a..adf8c6a 100644
--- a/vmtp/compute.py
+++ b/vmtp/compute.py
@@ -68,6 +68,17 @@ class Compute(object):
 
         return True
 
+    def delete_image(self, glance_client, img_name):
+        try:
+            print "Deleting image %s..." % img_name
+            img = glance_client.images.find(name=img_name)
+            glance_client.images.delete(img.id)
+        except Exception:
+            print "Failed to delete the image %s." % img_name
+            return False
+
+        return True
+
     # Remove keypair name from openstack if exists
     def remove_public_key(self, name):
         keypair_list = self.novaclient.keypairs.list()
diff --git a/vmtp/vmtp.py b/vmtp/vmtp.py
index 38b9227..bdcbc5e 100755
--- a/vmtp/vmtp.py
+++ b/vmtp/vmtp.py
@@ -145,6 +145,8 @@ class VmtpTest(object):
         self.image_instance = None
         self.flavor_type = None
         self.instance_access = None
+        self.glance_client = None
+        self.image_uploaded = False
         self.rescol = rescol
         self.config = config
         self.cred = cred
@@ -212,13 +214,14 @@ class VmtpTest(object):
                     keystone = keystoneclient.Client(**creds)
                     glance_endpoint = keystone.service_catalog.url_for(
                         service_type='image', endpoint_type='publicURL')
-                    glance_client = glanceclient.Client(
+                    self.glance_client = glanceclient.Client(
                         glance_endpoint, token=keystone.auth_token)
                     self.comp.upload_image_via_url(
-                        glance_client,
+                        self.glance_client,
                         self.config.image_name,
                         self.config.vm_image_url)
                     self.image_instance = self.comp.find_image(self.config.image_name)
+                    self.image_uploaded = True
                 else:
                     # Exit the pogram
                     print '%s: image to launch VM not found. ABORTING.' \
@@ -405,6 +408,8 @@ class VmtpTest(object):
         except ClientException:
             # May throw novaclient.exceptions.BadRequest if in use
             print('Security group in use: not deleted')
+        if self.image_uploaded and self.config.delete_image_after_run:
+            self.comp.delete_image(self.glance_client, self.config.image_name)
 
     def run(self):
         error_flag = False
@@ -719,6 +724,11 @@ def parse_opts_from_cli():
                         help='transport perf tool to use (default=nuttcp)',
                         metavar='<nuttcp|iperf>')
 
+    parser.add_argument('--availability_zone', dest='availability_zone',
+                        action='store',
+                        help='availability zone for running VMTP',
+                        metavar='<availability_zone>')
+
     # note there is a bug in argparse that causes an AssertionError
     # when the metavar is set to '[<az>:]<hostname>', hence had to insert a space
     parser.add_argument('--hypervisor', dest='hypervisors',
@@ -770,6 +780,11 @@ def parse_opts_from_cli():
                         help='Internal network name for OpenStack to hold data plane traffic',
                         metavar='<network_name>')
 
+    parser.add_argument('--delete-image-after-run', dest='delete_image_after_run',
+                        default=False,
+                        action='store_true',
+                        help='delete image that are uploaded by VMTP when tests are finished')
+
     parser.add_argument('--no-env', dest='no_env',
                         default=False,
                         action='store_true',
@@ -850,6 +865,9 @@ def merge_opts_to_configs(opts):
     config.vnic_type = opts.vnic_type
     config.hypervisors = opts.hypervisors
 
+    if opts.availability_zone:
+        config.availability_zone = opts.availability_zone
+
     # time to run each perf test in seconds
     if opts.time:
         config.time = int(opts.time)
@@ -923,6 +941,8 @@ def merge_opts_to_configs(opts):
     if opts.os_dataplane_network:
         config.os_dataplane_network = opts.os_dataplane_network
 
+    config.delete_image_after_run = opts.delete_image_after_run
+
     #####################################################
     # Set Ganglia server ip and port if the monitoring (-m)
     # option is enabled.