Merge "BugFix: Fix to negative dropped packets appearing in results."

[yardstick.git] / yardstick / benchmark / contexts / node.py

diff --git a/yardstick/benchmark/contexts/node.py b/yardstick/benchmark/contexts/node.py
index 35c6433..ffc82c8 100644 (file)
--- a/yardstick/benchmark/contexts/node.py
+++ b/yardstick/benchmark/contexts/node.py
@@ -17,12 +17,12 @@ import tempfile
 
 import six
 import pkg_resources
-import yaml
 
 from yardstick import ssh
 from yardstick.benchmark.contexts.base import Context
 from yardstick.common.constants import ANSIBLE_DIR, YARDSTICK_ROOT_PATH
 from yardstick.common.ansible_common import AnsibleCommon
+from yardstick.common.yaml_loader import yaml_load
 
 LOG = logging.getLogger(__name__)
 
@@ -55,7 +55,7 @@ class NodeContext(Context):
 
         with open(self.file_path) as stream:
             LOG.info("Parsing pod file: %s", self.file_path)
-            cfg = yaml.safe_load(stream)
+            cfg = yaml_load(stream)
         return cfg
 
     def init(self, attrs):
@@ -74,11 +74,11 @@ class NodeContext(Context):
 
         self.nodes.extend(cfg["nodes"])
         self.controllers.extend([node for node in cfg["nodes"]
-                                 if node["role"] == "Controller"])
+                                 if node.get("role") == "Controller"])
         self.computes.extend([node for node in cfg["nodes"]
-                              if node["role"] == "Compute"])
+                              if node.get("role") == "Compute"])
         self.baremetals.extend([node for node in cfg["nodes"]
-                                if node["role"] == "Baremetal"])
+                                if node.get("role") == "Baremetal"])
         LOG.debug("Nodes: %r", self.nodes)
         LOG.debug("Controllers: %r", self.controllers)
         LOG.debug("Computes: %r", self.computes)