test(apitester): rewrite json replacement logic to better handle arrays and existence (#4306)

Currently, attempting to replace properties within arrays on objects
that don't always have the same structure all the way down can result in
undesired behaviours such as the property being added to all objects
within arrays on the path, or only half of the potential matches being
replaced.

While the problem is a bit awkward to describe, the solution is not: now
we fully unroll a replacement path, and apply the replacement logic on
each resulting path.

i.e.with the object `{ "arr": [1, 2, 3] }`, we unroll the path `arr.#`
to `[]string{"arr.0", "arr.1", "arr.2"}`.

Author: G-Rath

tools/apitester/internal/jsonreplace/jsonreplace.go

@@ -25,66 +25,64 @@ func DoBytes(t *testing.T, json []byte, rules []Rule) []byte {
 	return json
 }
 
-// replaceJSONInput takes a gjson path and replaces all elements the path matches with the output of matcher
-func replaceJSONInput(t *testing.T, jsonInput []byte, path string, matcher func(toReplace gjson.Result) any) []byte {
+func expandArrayPaths(t *testing.T, jsonInput []byte, path string) []string {
 	t.Helper()
 
-	pathArray := []string{}
+	// split on the first intermediate #, if present
+	pathToArray, restOfPath, hasArrayPlaceholder := strings.Cut(path, ".#.")
+
+	// if there is no intermediate placeholder, check for (and cut) a terminal one
+	if !hasArrayPlaceholder {
+		pathToArray, hasArrayPlaceholder = strings.CutSuffix(path, ".#")
+	}
+
+	// if there are no array placeholders in the path, just return it
+	if !hasArrayPlaceholder {
+		return []string{path}
+	}
+
+	r := gjson.GetBytes(jsonInput, pathToArray)
 
-	// If there are more than 2 #, sjson cannot replace them directly. Iterate out all individual entries
-	if strings.Contains(path, "#") {
-		// Get the path ending with #
-		// E.g. results.#.packages.#.vulnerabilities => results.#.packages.#
-		numOfEntriesPath := path[:strings.LastIndex(path, "#")+1]
-		// This returns a potentially nested array of array lengths
-		numOfEntries := gjson.GetBytes(jsonInput, numOfEntriesPath)
+	// skip properties that are not arrays
+	if !r.IsArray() {
+		return []string{}
+	}
+
+	// if property exists and is actually an array, build out the path to each item
+	// within that array
+	paths := make([]string, 0, len(r.Array()))
+
+	for i := range r.Array() {
+		static := pathToArray + "." + strconv.Itoa(i)
 
-		// Use it to build up a list of concrete paths
-		buildSJSONPaths(t, &pathArray, path, numOfEntries)
-	} else {
-		pathArray = append(pathArray, path)
+		if restOfPath != "" {
+			static += "." + restOfPath
+		}
+		paths = append(paths, expandArrayPaths(t, jsonInput, static)...)
 	}
 
+	return paths
+}
+
+// replaceJSONInput takes a gjson path and replaces all elements the path matches with the output of matcher
+func replaceJSONInput(t *testing.T, jsonInput []byte, path string, replacer func(toReplace gjson.Result) any) []byte {
+	t.Helper()
+
 	var err error
 	json := jsonInput
-	for _, pathElem := range pathArray {
+	for _, pathElem := range expandArrayPaths(t, jsonInput, path) {
 		res := gjson.GetBytes(jsonInput, pathElem)
 
 		if !res.Exists() {
 			continue
 		}
 
-		json, err = sjson.SetBytesOptions(json, pathElem, matcher(res), &sjson.Options{Optimistic: true})
+		// optimistically replace the element, since we know at this point it does exist
+		json, err = sjson.SetBytesOptions(json, pathElem, replacer(res), &sjson.Options{Optimistic: true})
 		if err != nil {
 			t.Fatalf("failed to set element")
 		}
 	}
 
 	return json
 }
-
-func buildSJSONPaths(t *testing.T, pathToBuild *[]string, path string, structure gjson.Result) {
-	t.Helper()
-
-	if structure.IsArray() {
-		// More nesting to go
-		for i, res := range structure.Array() {
-			buildSJSONPaths(
-				t,
-				pathToBuild,
-				// Replace the first # with actual index
-				strings.Replace(path, "#", strconv.Itoa(i), 1),
-				res,
-			)
-		}
-	} else {
-		// Otherwise assume it is a number
-		if strings.Count(path, "#") != 1 {
-			t.Fatalf("programmer error: there should only be 1 # left")
-		}
-		for i2 := range int(structure.Int()) {
-			newPath := strings.Replace(path, "#", strconv.Itoa(i2), 1)
-			*pathToBuild = append(*pathToBuild, newPath)
-		}
-	}
-}