From e8d7f8e7c5189ab6a69462db9a0d75fa1c4cd259 Mon Sep 17 00:00:00 2001
From: debugtalk <mail@debugtalk.com>
Date: Fri, 16 Nov 2018 12:29:02 +0800
Subject: [PATCH] .env variables should be referenced with builtin ENV/environ
 function

---
 httprunner/exceptions.py     |  3 +++
 httprunner/parser.py         |  4 +++-
 httprunner/utils.py          | 19 +++++++++++++++++++
 tests/data/demo_testcase.yml |  1 +
 tests/test_loader.py         |  1 -
 tests/test_parser.py         | 10 ++++++++++
 6 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/httprunner/exceptions.py b/httprunner/exceptions.py
index 036137da..3db701ee 100644
--- a/httprunner/exceptions.py
+++ b/httprunner/exceptions.py
@@ -47,6 +47,9 @@ class FunctionNotFound(NotFoundError):
 class VariableNotFound(NotFoundError):
     pass
 
+class EnvNotFound(NotFoundError):
+    pass
+
 class ApiNotFound(NotFoundError):
     pass
 
diff --git a/httprunner/parser.py b/httprunner/parser.py
index ff68352a..d7fb2f01 100644
--- a/httprunner/parser.py
+++ b/httprunner/parser.py
@@ -415,7 +415,9 @@ def parse_string_functions(content, variables_mapping, functions_mapping):
 
         if func_name in ["parameterize", "P"]:
             from httprunner import loader
-            eval_value = loader.load_csv_file(*args, **kwargs)
+            eval_value = loader.load_csv_file(*args)
+        elif func_name in ["environ", "ENV"]:
+            eval_value = utils.get_os_environ(*args)
         else:
             func = get_mapping_function(func_name, functions_mapping)
             eval_value = func(*args, **kwargs)
diff --git a/httprunner/utils.py b/httprunner/utils.py
index d1ccc4a8..5fd35d8f 100644
--- a/httprunner/utils.py
+++ b/httprunner/utils.py
@@ -29,6 +29,25 @@ def unset_os_environ(variables_mapping):
         logger.log_debug("Unset OS environment variable: {}".format(variable))
 
 
+def get_os_environ(variable_name):
+    """ get value of environment variable.
+
+    Args:
+        variable_name(str): variable name
+
+    Returns:
+        value of environment variable.
+
+    Raises:
+        exceptions.EnvNotFound: If environment variable not found.
+
+    """
+    try:
+        return os.environ[variable_name]
+    except KeyError:
+        raise exceptions.EnvNotFound(variable_name)
+
+
 def query_json(json_content, query, delimiter='.'):
     """ Do an xpath-like query with json_content.
     @param (dict/list/string) json_content
diff --git a/tests/data/demo_testcase.yml b/tests/data/demo_testcase.yml
index 595185ef..e0f15764 100644
--- a/tests/data/demo_testcase.yml
+++ b/tests/data/demo_testcase.yml
@@ -3,6 +3,7 @@
     variables:
         - var_a: 0
         - var_c: "${sum_two(1, 2)}"
+        - PROJECT_KEY: ${ENV(PROJECT_KEY)}
     parameters:
         - "var_a-var_b":
             - [11, 21]
diff --git a/tests/test_loader.py b/tests/test_loader.py
index 2c6dcc49..91c279ed 100644
--- a/tests/test_loader.py
+++ b/tests/test_loader.py
@@ -27,7 +27,6 @@ class TestFileLoader(unittest.TestCase):
 
         os.remove(yaml_tmp_file)
 
-
     def test_load_json_file_file_format_error(self):
         json_tmp_file = "tests/data/tmp.json"
         # create empty file
diff --git a/tests/test_parser.py b/tests/test_parser.py
index 9dd7a6fd..b3a1ef17 100644
--- a/tests/test_parser.py
+++ b/tests/test_parser.py
@@ -432,8 +432,18 @@ class TestParser(unittest.TestCase):
         testcase_file_path = os.path.join(
             os.getcwd(), 'tests/data/demo_testcase.yml')
         testcases = loader.load_tests(testcase_file_path)
+        self.assertEqual(
+            testcases[0]["config"]["variables"][1]["var_c"],
+            "${sum_two(1, 2)}"
+        )
+        self.assertEqual(
+            testcases[0]["config"]["variables"][2]["PROJECT_KEY"],
+            "${ENV(PROJECT_KEY)}"
+        )
+
         parsed_testcases = parser.parse_tests(testcases)
         self.assertEqual(parsed_testcases[0]["config"]["variables"]["var_c"], 3)
+        self.assertEqual(parsed_testcases[0]["config"]["variables"]["PROJECT_KEY"], "ABCDEFGH")
         self.assertEqual(len(parsed_testcases), 2 * 2)
         self.assertEqual(
             parsed_testcases[0]["config"]["request"]["base_url"],