From 9367cac70609d58662ddcdcf4b04b6488c579746 Mon Sep 17 00:00:00 2001
From: debugtalk <mail@debugtalk.com>
Date: Sun, 7 Jun 2020 18:20:13 +0800
Subject: [PATCH] feat: step teardown hook

---
 examples/httpbin/hooks_test.py |  2 ++
 httprunner/make.py             | 11 +++++++++++
 httprunner/testcase.py         |  9 +++++++--
 3 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/examples/httpbin/hooks_test.py b/examples/httpbin/hooks_test.py
index eef0a46d..a10f629e 100644
--- a/examples/httpbin/hooks_test.py
+++ b/examples/httpbin/hooks_test.py
@@ -14,6 +14,7 @@ class TestCaseHooks(HttpRunner):
             .setup_hook("${setup_hook_add_kwargs($request)}")
             .setup_hook("${setup_hook_remove_kwargs($request)}")
             .get("/headers")
+            .teardown_hook("${teardown_hook_sleep_N_secs($response, 1)}")
             .validate()
             .assert_equal("status_code", 200)
             .assert_contained_by("body.headers.Host", "${get_httpbin_server()}")
@@ -21,6 +22,7 @@ class TestCaseHooks(HttpRunner):
         Step(
             RunRequest("alter response")
             .get("/headers")
+            .teardown_hook("${alter_response($response)}")
             .validate()
             .assert_equal("status_code", 200)
             .assert_equal("body.headers.Host", "httpbin.org")
diff --git a/httprunner/make.py b/httprunner/make.py
index a633efcd..967e8d74 100644
--- a/httprunner/make.py
+++ b/httprunner/make.py
@@ -246,6 +246,17 @@ def make_teststep_chain_style(teststep: Dict) -> Text:
         call_ref_testcase = f".call({testcase})"
         step_info += call_ref_testcase
 
+    if "teardown_hooks" in teststep:
+        teardown_hooks = teststep["teardown_hooks"]
+        for hook in teardown_hooks:
+            if isinstance(hook, Text):
+                step_info += f'.teardown_hook("{hook}")'
+            elif isinstance(hook, Dict) and len(hook) == 1:
+                assign_var_name, hook_content = list(hook.items())[0]
+                step_info += f'.teardown_hook("{hook}", "{assign_var_name}")'
+            else:
+                raise exceptions.TestCaseFormatError(f"Invalid teardown hook: {hook}")
+
     if "extract" in teststep:
         # request step
         step_info += ".extract()"
diff --git a/httprunner/testcase.py b/httprunner/testcase.py
index 42748381..5ceb4efc 100644
--- a/httprunner/testcase.py
+++ b/httprunner/testcase.py
@@ -265,8 +265,13 @@ class RequestWithOptionalArgs(object):
         self.__step_context.request.upload.update(file_info)
         return self
 
-    # def hooks(self):
-    #     pass
+    def teardown_hook(self, hook: Text, assign_var_name: Text = None) -> "RequestWithOptionalArgs":
+        if assign_var_name:
+            self.__step_context.teardown_hooks.append({assign_var_name: hook})
+        else:
+            self.__step_context.teardown_hooks.append(hook)
+
+        return self
 
     def extract(self) -> StepRequestExtraction:
         return StepRequestExtraction(self.__step_context)