refactor loading .env file

httprunner/loader.py

@@ -138,24 +138,31 @@ def load_folder_files(folder_path, recursive=True):
 
 
 def load_dot_env_file(path):
-    """ load .env file and set to os.environ
+    """ load .env file
     """
     if not path:
         path = os.path.join(os.getcwd(), ".env")
         if not os.path.isfile(path):
             logger.log_debug(".env file not exist: {}".format(path))
-            return
+            return {}
     else:
         if not os.path.isfile(path):
             raise exceptions.FileNotFound("env file not exist: {}".format(path))
 
     logger.log_info("Loading environment variables from {}".format(path))
+    env_variables_mapping = {}
     with io.open(path, 'r', encoding='utf-8') as fp:
         for line in fp:
-            variable, value = line.split("=")
-            variable = variable.strip()
-            os.environ[variable] = value.strip()
-            logger.log_debug("Loaded variable: {}".format(variable))
+            if "=" in line:
+                variable, value = line.split("=")
+            elif ":" in line:
+                variable, value = line.split(":")
+            else:
+                raise exceptions.FileFormatError(".env format error")
+
+            env_variables_mapping[variable.strip()] = value.strip()
+
+    return env_variables_mapping
 
 
 ###############################################################################

httprunner/task.py

@@ -206,7 +206,7 @@ class HttpRunner(object):
             - dot_env_path: .env file path
         """
         dot_env_path = kwargs.pop("dot_env_path", None)
-        loader.load_dot_env_file(dot_env_path)
+        utils.set_os_environ(loader.load_dot_env_file(dot_env_path))
 
         kwargs.setdefault("resultclass", HtmlTestResult)
         self.runner = unittest.TextTestRunner(**kwargs)

httprunner/utils.py

@@ -42,6 +42,14 @@ def remove_prefix(text, prefix):
     return text
 
 
+def set_os_environ(variables_mapping):
+    """ set variables mapping to os.environ
+    """
+    for variable in variables_mapping:
+        os.environ[variable] = variables_mapping[variable]
+        logger.log_debug("Loaded variable: {}".format(variable))
+
+
 def query_json(json_content, query, delimiter='.'):
     """ Do an xpath-like query with json_content.
     @param (dict/list/string) json_content

tests/test_loader.py

@@ -132,10 +132,9 @@ class TestFileLoader(unittest.TestCase):
         self.assertEqual([], files)
 
     def test_load_dot_env_file(self):
-        self.assertNotIn("PROJECT_KEY", os.environ)
-        loader.load_dot_env_file("tests/data/test.env")
-        self.assertIn("PROJECT_KEY", os.environ)
-        self.assertEqual(os.environ["UserName"], "debugtalk")
+        env_variables_mapping = loader.load_dot_env_file("tests/data/test.env")
+        self.assertIn("PROJECT_KEY", env_variables_mapping)
+        self.assertEqual(env_variables_mapping["UserName"], "debugtalk")
 
     def test_load_env_path_not_exist(self):
         with self.assertRaises(exceptions.FileNotFound):

tests/test_utils.py

@@ -17,6 +17,15 @@ class TestUtils(ApiServerUnittest):
             "/post/123"
         )
 
+    def test_set_os_environ(self):
+        self.assertNotIn("abc", os.environ)
+        variables_mapping = {
+            "abc": "123"
+        }
+        utils.set_os_environ(variables_mapping)
+        self.assertIn("abc", os.environ)
+        self.assertEqual(os.environ["abc"], "123")
+
     def test_query_json(self):
         json_content = {
             "ids": [1, 2, 3, 4],