From d8c1a66971f97213150aa33d9aad410b22804e7c Mon Sep 17 00:00:00 2001
From: debugtalk <mail@debugtalk.com>
Date: Mon, 14 May 2018 14:29:19 +0800
Subject: [PATCH] bugfix: modify request with setup_hooks

---
 httprunner/context.py | 12 ++++++++----
 httprunner/runner.py  | 14 +++++++-------
 tests/debugtalk.py    |  3 +++
 tests/test_runner.py  | 36 +++++++++++++++++++++++++++++++++---
 4 files changed, 51 insertions(+), 14 deletions(-)

diff --git a/httprunner/context.py b/httprunner/context.py
index a131a60f..8b260a3a 100644
--- a/httprunner/context.py
+++ b/httprunner/context.py
@@ -117,8 +117,13 @@ class Context(object):
             if level == "testset":
                 self.testset_shared_variables_mapping[variable_name] = variable_eval_value
 
-            self.testcase_variables_mapping[variable_name] = variable_eval_value
-            self.testcase_parser.update_binded_variables(self.testcase_variables_mapping)
+            self.bind_testcase_variable(variable_name, variable_eval_value)
+
+    def bind_testcase_variable(self, variable_name, variable_value):
+        """ bind and update testcase variables mapping
+        """
+        self.testcase_variables_mapping[variable_name] = variable_value
+        self.testcase_parser.update_binded_variables(self.testcase_variables_mapping)
 
     def bind_extracted_variables(self, variables):
         """ bind extracted variables to testset context
@@ -127,8 +132,7 @@ class Context(object):
         """
         for variable_name, value in variables.items():
             self.testset_shared_variables_mapping[variable_name] = value
-            self.testcase_variables_mapping[variable_name] = value
-            self.testcase_parser.update_binded_variables(self.testcase_variables_mapping)
+            self.bind_testcase_variable(variable_name, value)
 
     def __update_context_functions_config(self, level, config_mapping):
         """
diff --git a/httprunner/runner.py b/httprunner/runner.py
index 6ca7160e..8160f20b 100644
--- a/httprunner/runner.py
+++ b/httprunner/runner.py
@@ -142,7 +142,12 @@ class Runner(object):
 
         # prepare
         parsed_request = self.init_config(testcase_dict, level="testcase")
-        self.context.bind_variables({"request": parsed_request}, level="testcase")
+        self.context.bind_testcase_variable("request", parsed_request)
+
+        # setup hooks
+        setup_hooks = testcase_dict.get("setup_hooks", [])
+        setup_hooks.insert(0, "${setup_hook_prepare_kwargs($request)}")
+        self.do_hook_actions(setup_hooks)
 
         try:
             url = parsed_request.pop('url')
@@ -154,11 +159,6 @@ class Runner(object):
         logger.log_info("{method} {url}".format(method=method, url=url))
         logger.log_debug("request kwargs(raw): {kwargs}".format(kwargs=parsed_request))
 
-        # setup hooks
-        setup_hooks = testcase_dict.get("setup_hooks", [])
-        setup_hooks.insert(0, "${setup_hook_prepare_kwargs($request)}")
-        self.do_hook_actions(setup_hooks)
-
         # request
         resp = self.http_client_session.request(
             method,
@@ -170,7 +170,7 @@ class Runner(object):
         # teardown hooks
         teardown_hooks = testcase_dict.get("teardown_hooks", [])
         if teardown_hooks:
-            self.context.bind_variables({"response": resp}, level="testcase")
+            self.context.bind_testcase_variable("response", resp)
             self.do_hook_actions(teardown_hooks)
 
         # extract
diff --git a/tests/debugtalk.py b/tests/debugtalk.py
index 5ba51598..c896456d 100644
--- a/tests/debugtalk.py
+++ b/tests/debugtalk.py
@@ -84,3 +84,6 @@ def teardown_hook_sleep_N_secs(response, n_secs):
 
 def hook_print(msg):
     print(msg)
+
+def modify_headers_os_platform(request, os_platform):
+    request["headers"]["os_platform"] = os_platform
diff --git a/tests/test_runner.py b/tests/test_runner.py
index 35be9469..24874538 100644
--- a/tests/test_runner.py
+++ b/tests/test_runner.py
@@ -70,9 +70,6 @@ class TestRunner(ApiServerUnittest):
             self.test_runner.run_test(test)
 
     def test_run_testset_with_hooks(self):
-        testcase_file_path = os.path.join(
-            os.getcwd(), 'tests/httpbin/hooks.yml')
-
         start_time = time.time()
 
         config_dict = {
@@ -122,6 +119,39 @@ class TestRunner(ApiServerUnittest):
         # testset teardown hook has not been executed now
         self.assertLess(end_time - start_time, 1)
 
+    def test_run_testset_with_hooks_modify_request(self):
+        config_dict = {
+            "path": os.path.join(os.getcwd(), __file__),
+            "name": "basic test with httpbin",
+            "request": {
+                "base_url": "http://127.0.0.1:3458/"
+            }
+        }
+        test = {
+            "name": "modify request headers",
+            "request": {
+                "url": "/anything",
+                "method": "POST",
+                "headers": {
+                    "content-type": "application/json",
+                    "user_agent": "iOS/10.3",
+                    "os_platform": "ios"
+                },
+                "json": {
+                    "sign": "f1219719911caae89ccc301679857ebfda115ca2"
+                }
+            },
+            "setup_hooks": [
+                "${modify_headers_os_platform($request, android)}"
+            ],
+            "validate": [
+                {"check": "status_code", "expect": 200},
+                {"check": "content.headers.Os-Platform", "expect": "android"}
+            ]
+        }
+        test_runner = runner.Runner(config_dict)
+        test_runner.run_test(test)
+
     def test_run_httprunner_with_hooks(self):
         testcase_file_path = os.path.join(
             os.getcwd(), 'tests/httpbin/hooks.yml')