X-Git-Url: https://gerrit.opnfv.org/gerrit/gitweb?a=blobdiff_plain;f=yardstick%2Fbenchmark%2Fcontexts%2Fheat.py;h=f4c48f4a5163d52df842dc0eea0cf8233bbd499d;hb=347c2d26ca095ec8d3c7e189bbf1ded7e0a1a8e8;hp=ac85b6ffebb120ff2cebaefe6e451bff2f8ccf62;hpb=3fd6e3cb5bfee0c29e2e880b5ec019c4ccb16df5;p=yardstick.git

diff --git a/yardstick/benchmark/contexts/heat.py b/yardstick/benchmark/contexts/heat.py
index ac85b6ffe..f4c48f4a5 100644
--- a/yardstick/benchmark/contexts/heat.py
+++ b/yardstick/benchmark/contexts/heat.py
@@ -16,6 +16,7 @@ from collections import OrderedDict
 import ipaddress
 import pkg_resources
 
+from yardstick.benchmark import contexts
 from yardstick.benchmark.contexts.base import Context
 from yardstick.benchmark.contexts.model import Network
 from yardstick.benchmark.contexts.model import PlacementGroup, ServerGroup
@@ -47,7 +48,7 @@ def h_join(*args):
 class HeatContext(Context):
     """Class that represents a context in the logical model"""
 
-    __context_type__ = "Heat"
+    __context_type__ = contexts.CONTEXT_HEAT
 
     def __init__(self):
         self.stack = None
@@ -58,6 +59,7 @@ class HeatContext(Context):
         self.server_groups = []
         self.keypair_name = None
         self.secgroup_name = None
+        self.security_group = None
         self._server_map = {}
         self.attrs = {}
         self._image = None
@@ -117,8 +119,11 @@ class HeatContext(Context):
             return
 
         self.keypair_name = h_join(self.name, "key")
+
         self.secgroup_name = h_join(self.name, "secgroup")
 
+        self.security_group = attrs.get("security_group")
+
         self._image = attrs.get("image")
 
         self._flavor = attrs.get("flavor")
@@ -184,7 +189,7 @@ class HeatContext(Context):
                 self.flavors.add(flavor)
 
         template.add_keypair(self.keypair_name, self.name)
-        template.add_security_group(self.secgroup_name)
+        template.add_security_group(self.secgroup_name, self.security_group)
 
         for network in self.networks.values():
             # Using existing network
@@ -503,10 +508,12 @@ class HeatContext(Context):
         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))
         result = {
             "user": server.context.user,
             "pkey": pkey,
+            "key_filename": key_filename,
             "private_ip": server.private_ip,
             "interfaces": server.interfaces,
             "routing_table": self.generate_routing_table(server),