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 4952901..44cc92a 100644 (file)
--- a/yardstick/common/utils.py
+++ b/yardstick/common/utils.py
@@ -31,6 +31,7 @@ import six
 from flask import jsonify
 from six.moves import configparser
 from oslo_serialization import jsonutils
+from oslo_utils import encodeutils
 
 import yardstick
 
@@ -65,7 +66,7 @@ def itersubclasses(cls, _seen=None):
                 yield sub
 
 
-def import_modules_from_package(package):
+def import_modules_from_package(package, raise_exception=False):
     """Import modules given a package name
 
     :param: package - Full package name. For example: rally.deploy.engines
@@ -86,10 +87,27 @@ def import_modules_from_package(package):
         for module_name in missing_modules:
             try:
                 importlib.import_module(module_name)
-            except (ImportError, SyntaxError):
+            except (ImportError, SyntaxError) as exc:
+                if raise_exception:
+                    raise exc
                 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)
@@ -106,19 +124,23 @@ def remove_file(path):
             raise
 
 
-def execute_command(cmd):
+def execute_command(cmd, **kwargs):
     exec_msg = "Executing command: '%s'" % cmd
     logger.debug(exec_msg)
 
-    output = subprocess.check_output(cmd.split()).split(os.linesep)
-
-    return output
+    output = subprocess.check_output(cmd.split(), **kwargs)
+    return encodeutils.safe_decode(output, incoming='utf-8').split(os.linesep)
 
 
 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
@@ -410,3 +432,31 @@ def read_meminfo(ssh_client):
         output[match[0]] = match[1]
 
     return output
+
+
+def find_relative_file(path, task_path):
+    """
+    Find file in one of places: in abs of path or relative to a directory path,
+    in this order.
+
+    :param path:
+    :param task_path:
+    :return str: full path to file
+    """
+    # fixme: create schema to validate all fields have been provided
+    for lookup in [os.path.abspath(path), os.path.join(task_path, path)]:
+        try:
+            with open(lookup):
+                return lookup
+        except IOError:
+            pass
+    raise IOError(errno.ENOENT, 'Unable to find {} file'.format(path))
+
+
+def open_relative_file(path, task_path):
+    try:
+        return open(path)
+    except IOError as e:
+        if e.errno == errno.ENOENT:
+            return open(os.path.join(task_path, path))
+        raise