From 79be0b81a0e6b1dbdfe9b40e9c42b7aa9a6f2ee1 Mon Sep 17 00:00:00 2001
From: Adrian Zalewski <bardzusny@users.noreply.github.com>
Date: Fri, 10 Jun 2022 14:37:47 +0200
Subject: [PATCH] JSONPath scrubber (prototype).

---
 approvals.go                                  |  9 +++
 go.mod                                        |  2 +
 go.sum                                        |  2 +
 scrubber_json_path.go                         | 71 +++++++++++++++++++
 scrubber_test.go                              | 14 ++++
 ...SONBytesWithJSONPathScrubber.approved.json | 11 +++
 6 files changed, 109 insertions(+)
 create mode 100644 go.sum
 create mode 100644 scrubber_json_path.go
 create mode 100644 testdata/scrubber_test.TestVerifyJSONBytesWithJSONPathScrubber.approved.json

diff --git a/approvals.go b/approvals.go
index edcdf70..ce270ec 100644
--- a/approvals.go
+++ b/approvals.go
@@ -286,6 +286,15 @@ func (v verifyOptions) WithRegexScrubber(scrubber *regexp.Regexp, replacer strin
 	return v
 }
 
+// WithJSONPathScrubber allows you to 'scrub' dynamic data such as timestamps within your test input
+// and replace it with a static placeholder
+func (v verifyOptions) WithJSONPathScrubber(path string, replacer string) verifyOptions {
+	v.scrubbers = append(v.scrubbers, func(s string) string {
+		return scrubJSONPath(s, path, replacer)
+	})
+	return v
+}
+
 // WithExtension overrides the default file extension (.txt) for approval files.
 func (v verifyOptions) WithExtension(extension string) verifyOptions {
 	v.extWithDot = extension
diff --git a/go.mod b/go.mod
index c568e78..af31834 100644
--- a/go.mod
+++ b/go.mod
@@ -1,3 +1,5 @@
 module github.com/approvals/go-approval-tests
 
 go 1.12
+
+require github.com/bitly/go-simplejson v0.5.0 // indirect
diff --git a/go.sum b/go.sum
new file mode 100644
index 0000000..ac10d06
--- /dev/null
+++ b/go.sum
@@ -0,0 +1,2 @@
+github.com/bitly/go-simplejson v0.5.0 h1:6IH+V8/tVMab511d5bn4M7EwGXZf9Hj6i2xSwkNEM+Y=
+github.com/bitly/go-simplejson v0.5.0/go.mod h1:cXHtHw4XUPsvGaxgjIAn8PhEWG9NfngEKAMDJEczWVA=
diff --git a/scrubber_json_path.go b/scrubber_json_path.go
new file mode 100644
index 0000000..5a0c456
--- /dev/null
+++ b/scrubber_json_path.go
@@ -0,0 +1,71 @@
+package approvals
+
+import (
+	"strings"
+
+	"github.com/bitly/go-simplejson"
+)
+
+// converts received JSON string into a simpleJSON struct, scrubs the path with passed "scrubVal" using a separate util function, turns it into a string again
+// (pretty naive implementation)
+func scrubJSONPath(jsonStr string, path string, scrubVal interface{}) string {
+	js, err := simplejson.NewJson([]byte(jsonStr))
+	if err != nil {
+		panic(err)
+	}
+
+	scrubJSONPathSimpleJSON(js, path, scrubVal)
+
+	resS, err := js.EncodePretty()
+	if err != nil {
+		panic(err)
+	}
+	return string(resS)
+}
+
+// limited implementation of JSONPath for scrubbing values in simplejson structs:
+//
+// ScrubJSONPath(data, "foo.bar", "<bar_scrubbed>")
+// { foo: { bar: "123" } }                   -> { foo: { bar: "<bar_scrubbed>" }
+//
+// - ScrubJSONPath(data, "foo[*].bar", "<bar_arr_scrubbed>")
+// { foo: [{ bar: "123" }, { bar: "234" }] } -> { foo: [{ bar: "<bar_arr_scrubbed>" }, { bar: "<bar_arr_scrubbed>" }]) }
+func scrubJSONPathSimpleJSON(data *simplejson.Json, path string, scrubVal interface{}) {
+	// "foo[*].bar" -> ["foo", "[*]", "bar"]
+	var elems []string
+	for _, elDot := range strings.Split(path, ".") {
+		if elDot == "[*]" {
+			elems = append(elems, "[*]")
+		} else if strings.HasSuffix(elDot, "[*]") {
+			elems = append(elems, strings.TrimSuffix(elDot, "[*]"))
+			elems = append(elems, "[*]")
+		} else {
+			elems = append(elems, elDot)
+		}
+	}
+
+	// end of recursive iteration, set the value (if key is present) and return:
+	if len(elems) == 1 {
+		if _, found := data.CheckGet(elems[0]); found {
+			data.Set(elems[0], scrubVal)
+		}
+		return
+	}
+
+	leftPath := path
+	for _, el := range elems {
+		leftPath = strings.TrimPrefix(leftPath, el)
+		leftPath = strings.TrimPrefix(leftPath, ".")
+
+		kjkj // special key for array object keys:
+		if el == "[*]" {
+			for i := 0; i < len(data.MustArray()); i++ {
+				scrubJSONPathSimpleJSON(data.GetIndex(i), leftPath, scrubVal)
+			}
+		} else {
+			if _, found := data.CheckGet(el); found {
+				scrubJSONPathSimpleJSON(data.Get(el), leftPath, scrubVal)
+			}
+		}
+	}
+}
diff --git a/scrubber_test.go b/scrubber_test.go
index b94beb4..eca7207 100644
--- a/scrubber_test.go
+++ b/scrubber_test.go
@@ -86,3 +86,17 @@ func TestVerifyAllWithRegexScrubber(t *testing.T) {
 	xs := []string{"Christopher", "Llewellyn"}
 	approvals.VerifyAll(t, "uppercase", xs, func(x interface{}) string { return fmt.Sprintf("%s => %s", x, strings.ToUpper(x.(string))) }, opts)
 }
+
+func TestVerifyJSONBytesWithJSONPathScrubber(t *testing.T) {
+	opts := approvals.Options().
+		WithJSONPathScrubber("greeting", "Hi!!!!").
+		WithJSONPathScrubber("list[*].greeting", "Why hello!!!")
+
+	jb := []byte(`{
+		"greeting":"Hello",
+		"list":[
+			{ "greeting": "Hello 1" },
+			{ "greeting": "Hello 2"}]
+		}`)
+	approvals.VerifyJSONBytes(t, jb, opts)
+}
diff --git a/testdata/scrubber_test.TestVerifyJSONBytesWithJSONPathScrubber.approved.json b/testdata/scrubber_test.TestVerifyJSONBytesWithJSONPathScrubber.approved.json
new file mode 100644
index 0000000..8f73c59
--- /dev/null
+++ b/testdata/scrubber_test.TestVerifyJSONBytesWithJSONPathScrubber.approved.json
@@ -0,0 +1,11 @@
+{
+  "greeting": "Hi!!!!",
+  "list": [
+    {
+      "greeting": "Why hello!!!"
+    },
+    {
+      "greeting": "Why hello!!!"
+    }
+  ]
+}
\ No newline at end of file