From d9917084398e89de47b63de593f779d4af272953 Mon Sep 17 00:00:00 2001
From: debugtalk <mail@debugtalk.com>
Date: Sun, 7 Jun 2020 18:06:42 +0800
Subject: [PATCH] feat: step setup hook

---
 examples/httpbin/hooks_test.py |  4 +++-
 httprunner/make.py             | 11 +++++++++++
 httprunner/models.py           |  2 +-
 httprunner/runner.py           | 34 ++++++++++++++++++----------------
 httprunner/testcase.py         |  8 ++++++++
 5 files changed, 41 insertions(+), 18 deletions(-)

diff --git a/examples/httpbin/hooks_test.py b/examples/httpbin/hooks_test.py
index cec6ad42..eef0a46d 100644
--- a/examples/httpbin/hooks_test.py
+++ b/examples/httpbin/hooks_test.py
@@ -1,4 +1,4 @@
-# NOTE: Generated By HttpRunner v3.0.9
+# NOTE: Generated By HttpRunner v3.0.10
 # FROM: examples/httpbin/hooks.yml
 
 from httprunner import HttpRunner, Config, Step, RunRequest, RunTestCase
@@ -11,6 +11,8 @@ class TestCaseHooks(HttpRunner):
         Step(
             RunRequest("headers")
             .with_variables(**{"a": 123})
+            .setup_hook("${setup_hook_add_kwargs($request)}")
+            .setup_hook("${setup_hook_remove_kwargs($request)}")
             .get("/headers")
             .validate()
             .assert_equal("status_code", 200)
diff --git a/httprunner/make.py b/httprunner/make.py
index 8dcdbb5c..a633efcd 100644
--- a/httprunner/make.py
+++ b/httprunner/make.py
@@ -228,6 +228,17 @@ def make_teststep_chain_style(teststep: Dict) -> Text:
         variables = teststep["variables"]
         step_info += f".with_variables(**{variables})"
 
+    if "setup_hooks" in teststep:
+        setup_hooks = teststep["setup_hooks"]
+        for hook in setup_hooks:
+            if isinstance(hook, Text):
+                step_info += f'.setup_hook("{hook}")'
+            elif isinstance(hook, Dict) and len(hook) == 1:
+                assign_var_name, hook_content = list(hook.items())[0]
+                step_info += f'.setup_hook("{hook}", "{assign_var_name}")'
+            else:
+                raise exceptions.TestCaseFormatError(f"Invalid setup hook: {hook}")
+
     if teststep.get("request"):
         step_info += make_request_chain_style(teststep["request"])
     elif teststep.get("testcase"):
diff --git a/httprunner/models.py b/httprunner/models.py
index 4a470aed..61061bff 100644
--- a/httprunner/models.py
+++ b/httprunner/models.py
@@ -15,7 +15,7 @@ FunctionsMapping = Dict[Text, Callable]
 Headers = Dict[Text, Text]
 Cookies = Dict[Text, Text]
 Verify = bool
-Hooks = Union[List[Text], Dict[Text, Text]]
+Hooks = List[Union[Text, Dict[Text, Text]]]
 Export = List[Text]
 Validators = List[Dict]
 Env = Dict[Text, Any]
diff --git a/httprunner/runner.py b/httprunner/runner.py
index b1366dfa..7dc18242 100644
--- a/httprunner/runner.py
+++ b/httprunner/runner.py
@@ -93,12 +93,12 @@ class HttpRunner(object):
         """ call hook actions.
 
         Args:
-            hooks: each action in actions list maybe in two format.
+            hooks (list): each hook in hooks list maybe in two format.
 
-                format1 (dict): assignment, the value returned by hook function will be assigned to variable.
-                    {"var": "${func()}"}
-                format2 (str): only call hook functions.
+                format1 (str): only call hook functions.
                     ${func()}
+                format2 (dict): assignment, the value returned by hook function will be assigned to variable.
+                    {"var": "${func()}"}
 
             step_variables: current step variables to call hook, include two special variables
 
@@ -110,9 +110,18 @@ class HttpRunner(object):
         """
         logger.debug(f"call {hook_type} hook actions.")
 
-        if isinstance(hooks, Dict):
-            # format 1: {"var": "${func()}"}
-            for var_name, hook_content in hooks.items():
+        if not isinstance(hooks, List):
+            logger.error(f"Invalid hooks format: {hooks}")
+            return
+
+        for hook in hooks:
+            if isinstance(hook, Text):
+                # format 1: ["${func()}"]
+                logger.debug(f"call hook function: {hook}")
+                parse_data(hook, step_variables, self.__project_meta.functions)
+            elif isinstance(hook, Dict) and len(hook) == 1:
+                # format 2: {"var": "${func()}"}
+                var_name, hook_content = list(hook.items())[0]
                 hook_content_eval = parse_data(
                     hook_content, step_variables, self.__project_meta.functions
                 )
@@ -121,15 +130,8 @@ class HttpRunner(object):
                 )
                 logger.debug(f"assign variable: {var_name} = {hook_content_eval}")
                 step_variables[var_name] = hook_content_eval
-
-        elif isinstance(hooks, List):
-            # format 2: ["${func()}"]
-            for hook in hooks:
-                logger.debug(f"call hook function: {hook}")
-                parse_data(hook, step_variables, self.__project_meta.functions)
-
-        else:
-            logger.warning(f"Invalid hooks format: {hooks}")
+            else:
+                logger.error(f"Invalid hook format: {hook}")
 
     def __run_step_request(self, step: TStep) -> StepData:
         """run teststep: request"""
diff --git a/httprunner/testcase.py b/httprunner/testcase.py
index 91601df9..42748381 100644
--- a/httprunner/testcase.py
+++ b/httprunner/testcase.py
@@ -286,6 +286,14 @@ class RunRequest(object):
         self.__step_context.variables.update(variables)
         return self
 
+    def setup_hook(self, hook: Text, assign_var_name: Text = None) -> "RunRequest":
+        if assign_var_name:
+            self.__step_context.setup_hooks.append({assign_var_name: hook})
+        else:
+            self.__step_context.setup_hooks.append(hook)
+
+        return self
+
     def get(self, url: Text) -> RequestWithOptionalArgs:
         self.__step_context.request = TRequest(method=MethodEnum.GET, url=url)
         return RequestWithOptionalArgs(self.__step_context)