group file functions to utils.FileUtils

httprunner/utils.py

@@ -1,6 +1,7 @@
 # encoding: utf-8
 
 import copy
+import csv
 import hashlib
 import hmac
 import imp
@@ -14,6 +15,7 @@ import string
 import types
 from datetime import datetime
 
+import yaml
 from httprunner import exception, logger
 from httprunner.compat import OrderedDict, is_py2, is_py3
 from requests.structures import CaseInsensitiveDict
@@ -41,41 +43,126 @@ def remove_prefix(text, prefix):
         return text[len(prefix):]
     return text
 
-def load_folder_files(folder_path, recursive=True):
-    """ load folder path, return all files in list format.
-    @param
-        folder_path: specified folder path to load
-        recursive: if True, will load files recursively
-    """
-    if isinstance(folder_path, (list, set)):
-        files = []
-        for path in set(folder_path):
-            files.extend(load_folder_files(path, recursive))
 
-        return files
+class FileUtils(object):
 
-    if not os.path.exists(folder_path):
-        return []
+    def _check_format(file_path, content):
+        """ check testcase format if valid
+        """
+        if not content:
+            # testcase file content is empty
+            err_msg = u"Testcase file content is empty: {}".format(file_path)
+            logger.log_error(err_msg)
+            raise exception.FileFormatError(err_msg)
 
-    file_list = []
+        elif not isinstance(content, (list, dict)):
+            # testcase file content does not match testcase format
+            err_msg = u"Testcase file content format invalid: {}".format(file_path)
+            logger.log_error(err_msg)
+            raise exception.FileFormatError(err_msg)
 
-    for dirpath, dirnames, filenames in os.walk(folder_path):
-        filenames_list = []
+    def _load_yaml_file(yaml_file):
+        """ load yaml file and check file content format
+        """
+        with io.open(yaml_file, 'r', encoding='utf-8') as stream:
+            yaml_content = yaml.load(stream)
+            FileUtils._check_format(yaml_file, yaml_content)
+            return yaml_content
 
-        for filename in filenames:
-            if not filename.endswith(('.yml', '.yaml', '.json')):
-                continue
+    def _load_json_file(json_file):
+        """ load json file and check file content format
+        """
+        with io.open(json_file, encoding='utf-8') as data_file:
+            try:
+                json_content = json.load(data_file)
+            except exception.JSONDecodeError:
+                err_msg = u"JSONDecodeError: JSON file format error: {}".format(json_file)
+                logger.log_error(err_msg)
+                raise exception.FileFormatError(err_msg)
 
-            filenames_list.append(filename)
+            FileUtils._check_format(json_file, json_content)
+            return json_content
 
-        for filename in filenames_list:
-            file_path = os.path.join(dirpath, filename)
-            file_list.append(file_path)
+    def _load_csv_file(csv_file):
+        """ load csv file and check file content format
+        @param
+            csv_file: csv file path
+            e.g. csv file content:
+                username,password
+                test1,111111
+                test2,222222
+                test3,333333
+        @return
+            list of parameter, each parameter is in dict format
+            e.g.
+            [
+                {'username': 'test1', 'password': '111111'},
+                {'username': 'test2', 'password': '222222'},
+                {'username': 'test3', 'password': '333333'}
+            ]
+        """
+        csv_content_list = []
 
-        if not recursive:
-            break
+        with io.open(csv_file, encoding='utf-8') as csvfile:
+            reader = csv.DictReader(csvfile)
+            for row in reader:
+                csv_content_list.append(row)
+
+        return csv_content_list
+
+    def load_file(file_path):
+        if not os.path.isfile(file_path):
+            raise exception.FileNotFoundError("{} does not exist.".format(file_path))
+
+        file_suffix = os.path.splitext(file_path)[1].lower()
+        if file_suffix == '.json':
+            return FileUtils._load_json_file(file_path)
+        elif file_suffix in ['.yaml', '.yml']:
+            return FileUtils._load_yaml_file(file_path)
+        elif file_suffix == ".csv":
+            return FileUtils._load_csv_file(file_path)
+        else:
+            # '' or other suffix
+            err_msg = u"Unsupported file format: {}".format(file_path)
+            logger.log_warning(err_msg)
+            return []
+
+    def load_folder_files(folder_path, recursive=True):
+        """ load folder path, return all files in list format.
+        @param
+            folder_path: specified folder path to load
+            recursive: if True, will load files recursively
+        """
+        if isinstance(folder_path, (list, set)):
+            files = []
+            for path in set(folder_path):
+                files.extend(FileUtils.load_folder_files(path, recursive))
+
+            return files
+
+        if not os.path.exists(folder_path):
+            return []
+
+        file_list = []
+
+        for dirpath, dirnames, filenames in os.walk(folder_path):
+            filenames_list = []
+
+            for filename in filenames:
+                if not filename.endswith(('.yml', '.yaml', '.json')):
+                    continue
+
+                filenames_list.append(filename)
+
+            for filename in filenames_list:
+                file_path = os.path.join(dirpath, filename)
+                file_list.append(file_path)
+
+            if not recursive:
+                break
+
+        return file_list
 
-    return file_list
 
 def query_json(json_content, query, delimiter='.'):
     """ Do an xpath-like query with json_content.