From 155ea7dad83b689fdfde01eb43f725785308c3d9 Mon Sep 17 00:00:00 2001
From: debugtalk <mail@debugtalk.com>
Date: Sun, 7 Jun 2020 16:38:59 +0800
Subject: [PATCH] feat: implement step setup/teardown hooks

---
 docs/CHANGELOG.md    |  6 +++++
 httprunner/models.py | 10 ++++----
 httprunner/runner.py | 57 +++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 67 insertions(+), 6 deletions(-)

diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md
index e7a287da..4e9cc920 100644
--- a/docs/CHANGELOG.md
+++ b/docs/CHANGELOG.md
@@ -1,5 +1,11 @@
 # Release History
 
+## 3.0.10 (2020-06-07)
+
+**Added**
+
+- feat: implement step setup/teardown hooks
+
 ## 3.0.9 (2020-06-07)
 
 **Fixed**
diff --git a/httprunner/models.py b/httprunner/models.py
index dcfd40f2..4a470aed 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
-Hook = List[Text]
+Hooks = Union[List[Text], Dict[Text, Text]]
 Export = List[Text]
 Validators = List[Dict]
 Env = Dict[Text, Any]
@@ -37,8 +37,8 @@ class TConfig(BaseModel):
     base_url: BaseUrl = ""
     # Text: prepare variables in debugtalk.py, ${gen_variables()}
     variables: Union[VariablesMapping, Text] = {}
-    setup_hooks: Hook = []
-    teardown_hooks: Hook = []
+    # setup_hooks: Hooks = []
+    # teardown_hooks: Hooks = []
     export: Export = []
     path: Text = None
 
@@ -64,8 +64,8 @@ class TStep(BaseModel):
     request: Union[TRequest, None] = None
     testcase: Union[Text, Callable, None] = None
     variables: VariablesMapping = {}
-    setup_hooks: Hook = []
-    teardown_hooks: Hook = []
+    setup_hooks: Hooks = []
+    teardown_hooks: Hooks = []
     # used to extract request's response field
     extract: VariablesMapping = {}
     # used to export session variables from referenced testcase
diff --git a/httprunner/runner.py b/httprunner/runner.py
index c050deea..41cd4114 100644
--- a/httprunner/runner.py
+++ b/httprunner/runner.py
@@ -2,7 +2,7 @@ import os
 import time
 import uuid
 from datetime import datetime
-from typing import List, Dict, Text, NoReturn
+from typing import List, Dict, Text, NoReturn, Union
 
 try:
     import allure
@@ -31,6 +31,7 @@ from httprunner.models import (
     TestCaseInOut,
     ProjectMeta,
     TestCase,
+    Hooks,
 )
 
 
@@ -86,6 +87,50 @@ class HttpRunner(object):
         self.__export = export
         return self
 
+    def __call_hooks(
+        self, hooks: Hooks, step_variables: VariablesMapping, hook_type: Text,
+    ) -> NoReturn:
+        """ call hook actions.
+
+        Args:
+            hooks: each action in actions 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.
+                    ${func()}
+
+            step_variables: current step variables to call hook, include two special variables
+
+                request: parsed request dict
+                response: ResponseObject for current response
+
+            hook_type: setup/teardown
+
+        """
+        logger.debug(f"call {hook_type} hook actions.")
+
+        if isinstance(hooks, Dict):
+            # format 1: {"var": "${func()}"}
+            for var_name, hook_content in hooks.items():
+                hook_content_eval = parse_data(
+                    hook_content, step_variables, self.__project_meta.functions
+                )
+                logger.debug(
+                    f"call hook function: {hook_content}, got value: {hook_content_eval}"
+                )
+                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}")
+
     def __run_step_request(self, step: TStep) -> StepData:
         """run teststep: request"""
         step_data = StepData(name=step.name)
@@ -101,6 +146,11 @@ class HttpRunner(object):
             "HRUN-Request-ID",
             f"HRUN-{self.__case_id}-{str(int(time.time() * 1000))[-6:]}",
         )
+        step.variables["request"] = parsed_request_dict
+
+        # setup hooks
+        if step.setup_hooks:
+            self.__call_hooks(step.setup_hooks, step.variables, "setup")
 
         # prepare arguments
         method = parsed_request_dict.pop("method")
@@ -112,6 +162,11 @@ class HttpRunner(object):
         # request
         resp = self.__session.request(method, url, **parsed_request_dict)
         resp_obj = ResponseObject(resp)
+        step.variables["response"] = resp_obj
+
+        # teardown hooks
+        if step.teardown_hooks:
+            self.__call_hooks(step.teardown_hooks, step.variables, "teardown")
 
         def log_req_resp_details():
             err_msg = "\n{} DETAILED REQUEST & RESPONSE {}\n".format("*" * 32, "*" * 32)