Support using existing private key when using external heat template file

Currently, the private key is auto-generated from Yardstick.
This patch will allow user use their existing private key when
they use external heat template file.

JIRA: YARDSTICK-1447

Change-Id: I45a0ab4ac08e7ccedd770867ed43de92040b6a10
Signed-off-by: treyad <treyad@viosoft.com>

2 files changed, 80 insertions, 16 deletions

diff --git a/yardstick/benchmark/contexts/heat.py b/yardstick/benchmark/contexts/heat.py
index f4c48f4a5..403562c46 100644
--- a/yardstick/benchmark/contexts/heat.py
+++ b/yardstick/benchmark/contexts/heat.py
@@ -71,6 +71,7 @@ class HeatContext(Context):
         self.shade_client = None
         self.heat_timeout = None
         self.key_filename = None
+        self.yardstick_gen_key_file = True
         self.shade_client = None
         self.operator_client = None
         self.nodes = []
@@ -105,6 +106,14 @@ class HeatContext(Context):
 
         self.template_file = attrs.get("heat_template")
 
+        # try looking for external private key when using external heat template
+        if self.template_file is not None:
+            self.key_filename = attrs.get("key_filename", None)
+            if self.key_filename is not None:
+                # Disable key file generation if an external private key
+                # has been provided
+                self.yardstick_gen_key_file = False
+
         self.shade_client = openstack_utils.get_shade_client()
         self.operator_client = openstack_utils.get_shade_operator_client()
 
@@ -335,14 +344,16 @@ class HeatContext(Context):
         """deploys template into a stack using cloud"""
         LOG.info("Deploying context '%s' START", self.name)
 
-        self.key_filename = ''.join(
-            [consts.YARDSTICK_ROOT_PATH,
-             'yardstick/resources/files/yardstick_key-',
-             self.name])
+        # Check if there was no external private key provided
+        if self.key_filename is None:
+            self.key_filename = ''.join(
+                [consts.YARDSTICK_ROOT_PATH,
+                 'yardstick/resources/files/yardstick_key-',
+                 self.name])
         # Permissions may have changed since creation; this can be fixed. If we
         # overwrite the file, we lose future access to VMs using this key.
         # As long as the file exists, even if it is unreadable, keep it intact
-        if not os.path.exists(self.key_filename):
+        if self.yardstick_gen_key_file and not os.path.exists(self.key_filename):
             SSH.gen_keys(self.key_filename)
 
         heat_template = HeatTemplate(
@@ -442,12 +453,14 @@ class HeatContext(Context):
         }
 
     def _delete_key_file(self):
-        try:
-            utils.remove_file(self.key_filename)
-            utils.remove_file(self.key_filename + ".pub")
-        except OSError:
-            LOG.exception("There was an error removing the key file %s",
-                          self.key_filename)
+        # Only remove the key file if it has been generated by yardstick
+        if self.yardstick_gen_key_file:
+            try:
+                utils.remove_file(self.key_filename)
+                utils.remove_file(self.key_filename + ".pub")
+            except OSError:
+                LOG.exception("There was an error removing the key file %s",
+                              self.key_filename)
 
     def undeploy(self):
         """undeploys stack from cloud"""
@@ -505,11 +518,25 @@ class HeatContext(Context):
             if server is None:
                 return None
 
-        pkey = pkg_resources.resource_string(
-            'yardstick.resources',
-            h_join('files/yardstick_key', self.name)).decode('utf-8')
-        key_filename = pkg_resources.resource_filename('yardstick.resources',
-            h_join('files/yardstick_key', self.name))
+        # Get the pkey
+        if self.yardstick_gen_key_file:
+            pkey = pkg_resources.resource_string(
+                'yardstick.resources',
+                h_join('files/yardstick_key', self.name)).decode('utf-8')
+            key_filename = pkg_resources.resource_filename('yardstick.resources',
+                h_join('files/yardstick_key', self.name))
+        else:
+            # make sure the file exists before attempting to open it
+            if not os.path.exists(self.key_filename):
+                LOG.error("The key_filename provided %s does not exist!",
+                          self.key_filename)
+            else:
+                try:
+                    pkey = open(self.key_filename, 'r').read().decode('utf-8')
+                    key_filename = self.key_filename
+                except IOError:
+                    LOG.error("The key_filename provided (%s) is unreadable.",
+                              self.key_filename)
         result = {
             "user": server.context.user,
             "pkey": pkey,
diff --git a/yardstick/tests/unit/benchmark/contexts/test_heat.py b/yardstick/tests/unit/benchmark/contexts/test_heat.py
index 3ccae44c7..fbc6fd26d 100644
--- a/yardstick/tests/unit/benchmark/contexts/test_heat.py
+++ b/yardstick/tests/unit/benchmark/contexts/test_heat.py
@@ -81,6 +81,7 @@ class HeatContextTestCase(unittest.TestCase):
         self.assertIsNone(self.test_context.template_file)
         self.assertIsNone(self.test_context.heat_parameters)
         self.assertIsNone(self.test_context.key_filename)
+        self.assertTrue(self.test_context.yardstick_gen_key_file)
 
     @mock.patch.object(yaml_loader, 'read_yaml_file')
     @mock.patch('yardstick.benchmark.contexts.heat.PlacementGroup')
@@ -173,6 +174,23 @@ class HeatContextTestCase(unittest.TestCase):
         self.assertTrue(self.test_context._flags.no_setup)
         self.assertTrue(self.test_context._flags.no_teardown)
 
+    def test_init_key_filename(self):
+        attrs = {'name': 'foo',
+                 'file': 'pod.yaml',
+                 'task_id': '1234567890',
+                 'server_groups': {},
+                 'networks': {},
+                 'servers': {},
+                 'heat_template': "/root/clearwater.yaml",
+                 'key_filename': '/etc/yardstick/yardstick.pem'}
+
+        with mock.patch.object(openstack_utils, 'get_shade_client'), \
+             mock.patch.object(openstack_utils, 'get_shade_operator_client'):
+            self.test_context.init(attrs)
+
+        self.assertIsNotNone(self.test_context.key_filename)
+        self.assertFalse(self.test_context.yardstick_gen_key_file)
+
     @mock.patch('yardstick.benchmark.contexts.heat.HeatTemplate')
     def test__add_resources_to_template_no_servers(self, mock_template):
         self.test_context._name = 'ctx'
@@ -373,6 +391,25 @@ class HeatContextTestCase(unittest.TestCase):
         self.assertTrue(mock_manager.mock_calls.index(mock_call_gen_keys) <
                         mock_manager.mock_calls.index(mock_call_add_resources))
 
+    @mock.patch.object(heat, 'HeatTemplate')
+    @mock.patch.object(ssh.SSH, 'gen_keys')
+    @mock.patch.object(heat.HeatContext, '_create_new_stack')
+    def test_deploy_with_key_filename_provided(self, mock_create_new_stack,
+                                               mock_gen_keys, *args):
+        self.test_context._name = 'foo'
+        self.test_context._task_id = '1234567890'
+        self.test_context._name_task_id = '{}-{}'.format(
+            self.test_context._name, self.test_context._task_id[:8])
+        self.test_context.template_file = '/bar/baz/some-heat-file'
+        self.test_context.heat_parameters = {'image': 'cirros'}
+        self.test_context.yardstick_gen_key_file = False
+        self.test_context.key_filename = '/etc/yardstick/yardstick.pem'
+        self.test_context.get_neutron_info = mock.MagicMock()
+        self.test_context.deploy()
+
+        mock_create_new_stack.assert_called()
+        mock_gen_keys.assert_not_called()
+
     def test_check_for_context(self):
         pass
         # check that the context exists