From c0ebe2ee53a7b925b41e7a94ae77d26146842614 Mon Sep 17 00:00:00 2001
From: buyuxiang <347586493@qq.com>
Date: Sun, 20 Feb 2022 13:43:28 +0800
Subject: [PATCH] fix: extract data using jmespath

Change-Id: Icea3b8fa7e71dd610c19c91e21a259104ab2fe30
---
 docs/CHANGELOG.md  |  4 ++++
 examples/demo.json | 24 +++++++++++++++++++++++-
 examples/demo.yaml | 14 ++++++++++++++
 parser.go          | 13 +++++++++++--
 response.go        | 12 +++++++++++-
 5 files changed, 63 insertions(+), 4 deletions(-)

diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md
index e244f80c..dffc118b 100644
--- a/docs/CHANGELOG.md
+++ b/docs/CHANGELOG.md
@@ -1,5 +1,9 @@
 # Release History
 
+## v0.6.2 (2022-02-20)
+- change: json unmarshal to json.Number when parsing data
+- fix: incorrect data type when extracting data using jmespath
+
 ## v0.6.1 (2022-02-17)
 
 - fix: set request Content-Type for posting json only when not specified
diff --git a/examples/demo.json b/examples/demo.json
index 198bcbaa..2213444a 100644
--- a/examples/demo.json
+++ b/examples/demo.json
@@ -119,9 +119,13 @@
                 },
                 "body": {
                     "foo1": "$varFoo1",
-                    "foo2": "${max($a, $b)}"
+                    "foo2": "${max($a, $b)}",
+                    "time": "${get_timestamp()}"
                 }
             },
+            "extract": {
+                "varTime": "body.form.time"
+            },
             "validate": [
                 {
                     "check": "status_code",
@@ -142,6 +146,24 @@
                     "msg": "check args foo2"
                 }
             ]
+        },
+        {
+            "name": "final check",
+            "request": {
+                "method": "GET",
+                "url": "/get",
+                "params": {
+                    "time": "$varTime"
+                }
+            },
+            "validate": [
+                {
+                    "check": "body.args.time",
+                    "assert": "length_equals",
+                    "expect": 13,
+                    "msg": "check extracted var timestamp"
+                }
+            ]
         }
     ]
 }
\ No newline at end of file
diff --git a/examples/demo.yaml b/examples/demo.yaml
index ed6131dc..650c97c8 100644
--- a/examples/demo.yaml
+++ b/examples/demo.yaml
@@ -80,6 +80,9 @@ teststeps:
         body:
             foo1: $varFoo1
             foo2: ${max($a, $b)}
+            time: ${get_timestamp()}
+      extract:
+        varTime: body.form.time
       validate:
         - check: status_code
           assert: equals
@@ -93,3 +96,14 @@ teststeps:
           assert: equals
           expect: "12.3"
           msg: check args foo2
+    - name: final check
+      request:
+        method: GET
+        url: /get
+        params:
+            time: $varTime
+      validate:
+        - check: body.args.time
+          assert: length_equals
+          expect: 13
+          msg: check extracted var timestamp
diff --git a/parser.go b/parser.go
index 692faf75..ce2a8fe3 100644
--- a/parser.go
+++ b/parser.go
@@ -69,8 +69,7 @@ func (p *parser) parseData(raw interface{}, variablesMapping map[string]interfac
 	case reflect.String:
 		// json.Number
 		if rawValue, ok := raw.(json.Number); ok {
-			// use the same rule as json.Unmarshal (float64, for JSON numbers)
-			return rawValue.Float64()
+			return parseJSONNumber(rawValue)
 		}
 		// other string
 		value := rawValue.String()
@@ -109,6 +108,16 @@ func (p *parser) parseData(raw interface{}, variablesMapping map[string]interfac
 	}
 }
 
+func parseJSONNumber(raw json.Number) (interface{}, error) {
+	if strings.Contains(raw.String(), ".") {
+		// float64
+		return raw.Float64()
+	} else {
+		// int64
+		return raw.Int64()
+	}
+}
+
 const (
 	regexVariable     = `[a-zA-Z_]\w*`    // variable name should start with a letter or underscore
 	regexFunctionName = `[a-zA-Z_]\w*`    // function name should start with a letter or underscore
diff --git a/response.go b/response.go
index 29619160..9fc5fbaa 100644
--- a/response.go
+++ b/response.go
@@ -1,6 +1,7 @@
 package hrp
 
 import (
+	"bytes"
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
@@ -53,7 +54,9 @@ func newResponseObject(t *testing.T, parser *parser, resp *http.Response) (*resp
 	// convert respObjMeta to interface{}
 	respObjMetaBytes, _ := json.Marshal(respObjMeta)
 	var data interface{}
-	if err := json.Unmarshal(respObjMetaBytes, &data); err != nil {
+	decoder := json.NewDecoder(bytes.NewReader(respObjMetaBytes))
+	decoder.UseNumber()
+	if err := decoder.Decode(&data); err != nil {
 		log.Error().
 			Str("respObjMeta", string(respObjMetaBytes)).
 			Err(err).
@@ -167,5 +170,12 @@ func (v *responseObject) searchJmespath(expr string) interface{} {
 		log.Error().Str("expr", expr).Err(err).Msg("search jmespath failed")
 		return expr // jmespath not found, return the expression
 	}
+	if number, ok := checkValue.(json.Number); ok {
+		checkNumber, err := parseJSONNumber(number)
+		if err != nil {
+			log.Error().Interface("json number", number).Err(err).Msg("convert json number failed")
+		}
+		return checkNumber
+	}
 	return checkValue
 }