From 777c144813d5e16fc16d23993f6a2c8b90d9aa25 Mon Sep 17 00:00:00 2001
From: debugtalk <mail@debugtalk.com>
Date: Tue, 28 Aug 2018 14:35:26 +0800
Subject: [PATCH] add --dot-env-path back

---
 httprunner/api.py    |  4 +++-
 httprunner/cli.py    |  5 ++++-
 httprunner/loader.py | 15 +++++++++++----
 tests/data/test.env  |  3 +++
 tests/test_loader.py | 10 ++++++++++
 5 files changed, 31 insertions(+), 6 deletions(-)
 create mode 100644 tests/data/test.env

diff --git a/httprunner/api.py b/httprunner/api.py
index 2e87c5fe..b3356732 100644
--- a/httprunner/api.py
+++ b/httprunner/api.py
@@ -19,6 +19,7 @@ class HttpRunner(object):
             resultclass (class): HtmlTestResult or TextTestResult
             failfast (bool): False/True, stop the test run on the first error or failure.
             http_client_session (instance): requests.Session(), or locust.client.Session() instance.
+            dot_env_path (str): .env file path.
 
         Attributes:
             project_mapping (dict): save project loaded api/testcases, environments and debugtalk.py module.
@@ -33,8 +34,9 @@ class HttpRunner(object):
                 }
 
         """
+        loader.dot_env_path = kwargs.pop("dot_env_path", None)
+        self.http_client_session = kwargs.pop("http_client_session", None)
         self.kwargs = kwargs
-        self.http_client_session = self.kwargs.pop("http_client_session", None)
 
     def load_tests(self, path_or_testcases):
         """ load testcases, extend and merge with api/testcase definitions.
diff --git a/httprunner/cli.py b/httprunner/cli.py
index fc15ab5b..ec645471 100644
--- a/httprunner/cli.py
+++ b/httprunner/cli.py
@@ -39,6 +39,9 @@ def main_hrun():
     parser.add_argument(
         '--log-file',
         help="Write logs to specified file path.")
+    parser.add_argument(
+        '--dot-env-path',
+        help="Specify .env file path, which is useful for keeping sensitive data.")
     parser.add_argument(
         '--failfast', action='store_true', default=False,
         help="Stop the test run on the first error or failure.")
@@ -75,7 +78,7 @@ def main_hrun():
         create_scaffold(project_path)
         exit(0)
 
-    runner = HttpRunner(failfast=args.failfast).run(args.testset_paths)
+    runner = HttpRunner(failfast=args.failfast, dot_env_path=args.dot_env_path).run(args.testset_paths)
 
     if not args.no_html_report:
         runner.gen_html_report(
diff --git a/httprunner/loader.py b/httprunner/loader.py
index e17c46d4..ae6e5a85 100644
--- a/httprunner/loader.py
+++ b/httprunner/loader.py
@@ -24,6 +24,8 @@ project_mapping = {
 """ dict: save project loaded api/testcases definitions, environments and debugtalk.py module.
 """
 
+dot_env_path = None
+
 testcases_cache_mapping = {}
 project_working_directory = os.getcwd()
 
@@ -161,7 +163,8 @@ def load_folder_files(folder_path, recursive=True):
 
 
 def load_dot_env_file():
-    """ load .env file, .env file should be located in project working directory.
+    """ load .env file, .env file should be located in project working directory by default.
+        If dot_env_path is specified, it will be loaded instead.
 
     Returns:
         dict: environment variables mapping
@@ -176,10 +179,14 @@ def load_dot_env_file():
         exceptions.FileFormatError: If env file format is invalid.
 
     """
-    path = os.path.join(project_working_directory, ".env")
+    path = dot_env_path or os.path.join(project_working_directory, ".env")
     if not os.path.isfile(path):
-        logger.log_debug(".env file not exist in : {}".format(project_working_directory))
-        return {}
+        if dot_env_path:
+            logger.log_error(".env file not exist: {}".format(dot_env_path))
+            sys.exit(1)
+        else:
+            logger.log_debug(".env file not exist in: {}".format(project_working_directory))
+            return {}
 
     logger.log_info("Loading environment variables from {}".format(path))
     env_variables_mapping = {}
diff --git a/tests/data/test.env b/tests/data/test.env
new file mode 100644
index 00000000..09816d83
--- /dev/null
+++ b/tests/data/test.env
@@ -0,0 +1,3 @@
+UserName=test
+Password=654321
+PROJECT_KEY=AAABBBCCC
\ No newline at end of file
diff --git a/tests/test_loader.py b/tests/test_loader.py
index b6f08745..18e67d86 100644
--- a/tests/test_loader.py
+++ b/tests/test_loader.py
@@ -140,6 +140,16 @@ class TestFileLoader(unittest.TestCase):
         self.assertIn("PROJECT_KEY", env_variables_mapping)
         self.assertEqual(env_variables_mapping["UserName"], "debugtalk")
 
+    def test_load_custom_dot_env_file(self):
+        loader.project_working_directory = os.path.join(
+            os.getcwd(), "tests",
+        )
+        loader.dot_env_path = "tests/data/test.env"
+        env_variables_mapping = loader.load_dot_env_file()
+        self.assertIn("PROJECT_KEY", env_variables_mapping)
+        self.assertEqual(env_variables_mapping["UserName"], "test")
+        loader.dot_env_path = None
+
     def test_load_env_path_not_exist(self):
         loader.project_working_directory = os.path.join(
             os.getcwd(), "tests", "data",