From fcb1fb1575d4b2ed7bf2f53e36183fa5679a490d Mon Sep 17 00:00:00 2001
From: buyuxiang <347586493@qq.com>
Date: Tue, 12 Apr 2022 20:07:19 +0800
Subject: [PATCH] feat: parse checkItem before searchField

---
 hrp/response.go           | 23 +++++++++++++++++------
 hrp/response_test.go      | 24 ++++++++++++++++++++++++
 hrp/tests/request_test.go |  5 +++--
 3 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/hrp/response.go b/hrp/response.go
index 1396f2e2..68f70b75 100644
--- a/hrp/response.go
+++ b/hrp/response.go
@@ -91,7 +91,7 @@ type responseObject struct {
 
 const textExtractorSubRegexp string = `(.*)`
 
-func (v *responseObject) extractField(value string) interface{} {
+func (v *responseObject) searchField(value string) interface{} {
 	var result interface{}
 	if strings.Contains(value, textExtractorSubRegexp) {
 		result = v.searchRegexp(value)
@@ -108,7 +108,7 @@ func (v *responseObject) Extract(extractors map[string]string) map[string]interf
 
 	extractMapping := make(map[string]interface{})
 	for key, value := range extractors {
-		extractedValue := v.extractField(value)
+		extractedValue := v.searchField(value)
 		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
@@ -127,14 +127,16 @@ func (v *responseObject) Validate(iValidators []interface{}, variablesMapping ma
 		checkItem := validator.Check
 		var checkValue interface{}
 		if strings.Contains(checkItem, "$") {
-			// reference variable
+			// parse reference variables in checkItem
 			checkValue, err = v.parser.Parse(checkItem, variablesMapping)
 			if err != nil {
 				return err
 			}
 		} else {
-			// regExp or jmesPath
-			checkValue = v.extractField(checkItem)
+			checkValue = checkItem
+		}
+		if searchCheckValue, ok := checkValue.(string); ok && checkSearchField(searchCheckValue) {
+			checkValue = v.searchField(searchCheckValue)
 		}
 
 		// get assert method
@@ -187,6 +189,15 @@ func (v *responseObject) Validate(iValidators []interface{}, variablesMapping ma
 	return nil
 }
 
+func checkSearchField(expr string) bool {
+	return strings.Contains(expr, "proto") ||
+		strings.Contains(expr, "status_code") ||
+		strings.Contains(expr, "headers") ||
+		strings.Contains(expr, "cookies") ||
+		strings.Contains(expr, "body") ||
+		strings.Contains(expr, textExtractorSubRegexp)
+}
+
 func (v *responseObject) searchJmespath(expr string) interface{} {
 	checkValue, err := jmespath.Search(expr, v.respObjMeta)
 	if err != nil {
@@ -220,7 +231,7 @@ func (v *responseObject) searchRegexp(expr string) interface{} {
 		return expr
 	}
 	match := regexpCompile.FindStringSubmatch(bodyStr)
-	if match != nil || len(match) > 1 {
+	if len(match) > 1 {
 		return match[1] //return first matched result in parentheses
 	}
 	log.Error().Str("expr", expr).Msg("search regexp failed")
diff --git a/hrp/response_test.go b/hrp/response_test.go
index cfd3143d..49a32aa2 100644
--- a/hrp/response_test.go
+++ b/hrp/response_test.go
@@ -9,6 +9,30 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+func TestSearchJmespath(t *testing.T) {
+	testText := `{"a": {"b": "foo"}, "c": "bar", "d": {"e": [{"f": "foo"}, {"f": "bar"}]}}`
+	testData := []struct {
+		raw      string
+		expected string
+	}{
+		{"body.a.b", "foo"},
+		{"body.c", "bar"},
+		{"body.d.e[0].f", "foo"},
+		{"body.d.e[1].f", "bar"},
+	}
+	resp := http.Response{}
+	resp.Body = io.NopCloser(strings.NewReader(testText))
+	respObj, err := newResponseObject(t, newParser(), &resp)
+	if err != nil {
+		t.Fail()
+	}
+	for _, data := range testData {
+		if !assert.Equal(t, data.expected, respObj.searchJmespath(data.raw)) {
+			t.Fail()
+		}
+	}
+}
+
 func TestSearchRegexp(t *testing.T) {
 	testText := `hrp aims to be a one-stop solution for HTTP(S) testing, covering API testing, load testing and digital experience monitoring (DEM).`
 	testData := []struct {
diff --git a/hrp/tests/request_test.go b/hrp/tests/request_test.go
index 922734a1..59f228df 100644
--- a/hrp/tests/request_test.go
+++ b/hrp/tests/request_test.go
@@ -13,6 +13,7 @@ func TestCaseBasicRequest(t *testing.T) {
 			SetVerifySSL(false),
 		TestSteps: []hrp.IStep{
 			hrp.NewStep("get with params").
+				WithVariables(map[string]interface{}{"headers_key": "\"Content-Type\"", "body_key": "foo1"}).
 				GET("/get").
 				WithParams(map[string]interface{}{"foo1": "bar1", "foo2": "bar2"}).
 				WithHeaders(map[string]string{
@@ -20,8 +21,8 @@ func TestCaseBasicRequest(t *testing.T) {
 				}).
 				Validate().
 				AssertEqual("status_code", 200, "check status code").
-				AssertEqual("headers.\"Content-Type\"", "application/json; charset=utf-8", "check header Content-Type").
-				AssertEqual("body.args.foo1", "bar1", "check args foo1").
+				AssertEqual("headers.$headers_key", "application/json; charset=utf-8", "check header Content-Type").
+				AssertEqual("body.args.$body_key", "bar1", "check args foo1").
 				AssertEqual("body.args.foo2", "bar2", "check args foo2"),
 			hrp.NewStep("post raw text").
 				POST("/post").