Merge "Bugfix: Can't get image list in API"

[yardstick.git] / yardstick / common / utils.py

diff --git a/yardstick/common/utils.py b/yardstick/common/utils.py
index 3f2d546..44cc92a 100644 (file)
--- a/yardstick/common/utils.py
+++ b/yardstick/common/utils.py
@@ -93,6 +93,21 @@ def import_modules_from_package(package, raise_exception=False):
                 logger.exception('Unable to import module %s', module_name)
 
 
+NON_NONE_DEFAULT = object()
+
+
+def get_key_with_default(data, key, default=NON_NONE_DEFAULT):
+    value = data.get(key, default)
+    if value is NON_NONE_DEFAULT:
+        raise KeyError(key)
+    return value
+
+
+def make_dict_from_map(data, key_map):
+    return {dest_key: get_key_with_default(data, src_key, default)
+            for dest_key, (src_key, default) in key_map.items()}
+
+
 def makedirs(d):
     try:
         os.makedirs(d)
@@ -121,6 +136,11 @@ def source_env(env_file):
     p = subprocess.Popen(". %s; env" % env_file, stdout=subprocess.PIPE,
                          shell=True)
     output = p.communicate()[0]
+
+    # sometimes output type would be binary_type, and it don't have splitlines
+    # method, so we need to decode
+    if isinstance(output, six.binary_type):
+        output = encodeutils.safe_decode(output)
     env = dict(line.split('=', 1) for line in output.splitlines() if '=' in line)
     os.environ.update(env)
     return env