From f986a2298aaa938ed9f995d235c7dbd97a496db2 Mon Sep 17 00:00:00 2001
From: buyuxiang <347586493@qq.com>
Date: Fri, 24 Jun 2022 17:46:22 +0800
Subject: [PATCH] support variable reference during extraction

---
 docs/CHANGELOG.md         |  4 +++-
 hrp/response.go           | 41 +++++++++++++++++++--------------------
 hrp/step_request.go       |  2 +-
 hrp/step_websocket.go     |  2 +-
 hrp/tests/extract_test.go |  3 ++-
 5 files changed, 27 insertions(+), 25 deletions(-)

diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md
index 0bda2b7d..62ff1977 100644
--- a/docs/CHANGELOG.md
+++ b/docs/CHANGELOG.md
@@ -1,6 +1,6 @@
 # Release History
 
-## v4.1.5 (2022-06-26)
+## v4.1.5 (2022-06-27)
 
 **go version**
 
@@ -12,6 +12,8 @@
 - fix: failed to load json/data content in api reference
 - fix: failed to convert postman collection containing multipart/form-data requests to pytest
 - fix: only get the first parameter in referenced testcase
+- fix: support variable reference during extraction
+- fix: simplify jmespath compatibility conversion
 
 **python version**
 
diff --git a/hrp/response.go b/hrp/response.go
index 33b363cb..43c2d80b 100644
--- a/hrp/response.go
+++ b/hrp/response.go
@@ -123,24 +123,35 @@ type responseObject struct {
 
 const textExtractorSubRegexp string = `(.*)`
 
-func (v *responseObject) searchField(value string) interface{} {
-	var result interface{}
-	if strings.Contains(value, textExtractorSubRegexp) {
-		result = v.searchRegexp(value)
-	} else {
-		result = v.searchJmespath(value)
+func (v *responseObject) searchField(field string, variablesMapping map[string]interface{}) interface{} {
+	var result interface{} = field
+	if strings.Contains(field, "$") {
+		// parse reference variables in field before search
+		var err error
+		result, err = v.parser.Parse(field, variablesMapping)
+		if err != nil {
+			log.Error().Str("filed name", field).Err(err).Msg("fail to parse field before search")
+		}
+	}
+	// search field using jmespath or regex if parsed field is still string and contains specified fieldTags
+	if parsedField, ok := result.(string); ok && checkSearchField(parsedField) {
+		if strings.Contains(field, textExtractorSubRegexp) {
+			result = v.searchRegexp(parsedField)
+		} else {
+			result = v.searchJmespath(parsedField)
+		}
 	}
 	return result
 }
 
-func (v *responseObject) Extract(extractors map[string]string) map[string]interface{} {
+func (v *responseObject) Extract(extractors map[string]string, variablesMapping map[string]interface{}) map[string]interface{} {
 	if extractors == nil {
 		return nil
 	}
 
 	extractMapping := make(map[string]interface{})
 	for key, value := range extractors {
-		extractedValue := v.searchField(value)
+		extractedValue := v.searchField(value, variablesMapping)
 		log.Info().Str("from", value).Interface("value", extractedValue).Msg("extract value")
 		log.Info().Str("variable", key).Interface("value", extractedValue).Msg("set variable")
 		extractMapping[key] = extractedValue
@@ -157,19 +168,7 @@ func (v *responseObject) Validate(iValidators []interface{}, variablesMapping ma
 		}
 		// parse check value
 		checkItem := validator.Check
-		var checkValue interface{}
-		if strings.Contains(checkItem, "$") {
-			// parse reference variables in checkItem
-			checkValue, err = v.parser.Parse(checkItem, variablesMapping)
-			if err != nil {
-				return err
-			}
-		} else {
-			checkValue = checkItem
-		}
-		if searchCheckValue, ok := checkValue.(string); ok && checkSearchField(searchCheckValue) {
-			checkValue = v.searchField(searchCheckValue)
-		}
+		checkValue := v.searchField(checkItem, variablesMapping)
 
 		// get assert method
 		assertMethod := validator.Assert
diff --git a/hrp/step_request.go b/hrp/step_request.go
index d2ce333a..29a86578 100644
--- a/hrp/step_request.go
+++ b/hrp/step_request.go
@@ -419,7 +419,7 @@ func runStepRequest(r *SessionRunner, step *TStep) (stepResult *StepResult, err
 
 	// extract variables from response
 	extractors := step.Extract
-	extractMapping := respObj.Extract(extractors)
+	extractMapping := respObj.Extract(extractors, stepVariables)
 	stepResult.ExportVars = extractMapping
 
 	// override step variables with extracted variables
diff --git a/hrp/step_websocket.go b/hrp/step_websocket.go
index f3a788fd..34cc226c 100644
--- a/hrp/step_websocket.go
+++ b/hrp/step_websocket.go
@@ -365,7 +365,7 @@ func runStepWebSocket(r *SessionRunner, step *TStep) (stepResult *StepResult, er
 
 		// extract variables from response
 		extractors := step.Extract
-		extractMapping := respObj.Extract(extractors)
+		extractMapping := respObj.Extract(extractors, stepVariables)
 		stepResult.ExportVars = extractMapping
 
 		// override step variables with extracted variables
diff --git a/hrp/tests/extract_test.go b/hrp/tests/extract_test.go
index df8d2d18..9c2a8848 100644
--- a/hrp/tests/extract_test.go
+++ b/hrp/tests/extract_test.go
@@ -18,6 +18,7 @@ func TestCaseExtractStepSingle(t *testing.T) {
 					"var1":               "bar1",
 					"agent":              "HttpRunnerPlus",
 					"expectedStatusCode": 200,
+					"jmespathFoo1":       "body.args.foo1",
 				}).
 				GET("/get").
 				WithParams(map[string]interface{}{"foo1": "$var1", "foo2": "bar2"}).
@@ -25,7 +26,7 @@ func TestCaseExtractStepSingle(t *testing.T) {
 				Extract().
 				WithJmesPath("status_code", "statusCode").
 				WithJmesPath("headers.\"Content-Type\"", "contentType").
-				WithJmesPath("body.args.foo1", "varFoo1").
+				WithJmesPath("$jmespathFoo1", "varFoo1").
 				Validate().
 				AssertEqual("$statusCode", "$expectedStatusCode", "check status code").                      // assert with extracted variable from current step
 				AssertEqual("$contentType", "application/json; charset=utf-8", "check header Content-Type"). // assert with extracted variable from current step