From f80ff931db2bc4087e730e0155b410b426df96de Mon Sep 17 00:00:00 2001
From: cheetah <185541429@qq.com>
Date: Mon, 9 Jan 2023 23:48:52 +0800
Subject: [PATCH] =?UTF-8?q?setup=20hooks=20=E5=85=BC=E5=AE=B9v3=E7=89=88?=
 =?UTF-8?q?=E6=9C=AC=E5=90=8C=E6=A0=B7=E7=9A=84=E5=8A=9F=E8=83=BD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 hrp/step_request.go | 31 +++++++++++++++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

diff --git a/hrp/step_request.go b/hrp/step_request.go
index 12a8ffbb..abf9755f 100644
--- a/hrp/step_request.go
+++ b/hrp/step_request.go
@@ -337,13 +337,34 @@ func runStepRequest(r *SessionRunner, step *TStep) (stepResult *StepResult, err
 	// add request object to step variables, could be used in setup hooks
 	stepVariables["hrp_step_name"] = step.Name
 	stepVariables["hrp_step_request"] = rb.requestMap
+	stepVariables["request"] = rb.requestMap
 
 	// deal with setup hooks
 	for _, setupHook := range step.SetupHooks {
-		_, err = parser.Parse(setupHook, stepVariables)
+		req, err := parser.Parse(setupHook, stepVariables)
 		if err != nil {
 			return stepResult, errors.Wrap(err, "run setup hooks failed")
 		}
+		reqMap, ok := req.(map[string]interface{})
+		if ok && reqMap != nil {
+			rb.requestMap = reqMap
+			stepVariables["request"] = reqMap
+		}
+	}
+	if len(step.SetupHooks) > 0 {
+		requestBody, ok := rb.requestMap["body"].(map[string]interface{})
+		if ok {
+			body, err := json.Marshal(requestBody)
+			if err == nil {
+				rb.req.Body = io.NopCloser(bytes.NewReader(body))
+				rb.req.ContentLength = int64(len(body))
+			}
+		}
+		headers, ok := rb.requestMap["headers"].(map[string]string)
+		rb.req.Header = map[string][]string{}
+		for key, value := range headers {
+			rb.req.Header.Set(key, value)
+		}
 	}
 
 	// log & print request
@@ -414,13 +435,19 @@ func runStepRequest(r *SessionRunner, step *TStep) (stepResult *StepResult, err
 
 	// add response object to step variables, could be used in teardown hooks
 	stepVariables["hrp_step_response"] = respObj.respObjMeta
+	stepVariables["response"] = respObj.respObjMeta
 
 	// deal with teardown hooks
 	for _, teardownHook := range step.TeardownHooks {
-		_, err = parser.Parse(teardownHook, stepVariables)
+		res, err := parser.Parse(teardownHook, stepVariables)
 		if err != nil {
 			return stepResult, errors.Wrap(err, "run teardown hooks failed")
 		}
+		resMpa, ok := res.(map[string]interface{})
+		if ok {
+			stepVariables["response"] = resMpa
+			respObj.respObjMeta = resMpa
+		}
 	}
 
 	sessionData.ReqResps.Request = rb.requestMap