feat: add WithReplacement except opt to snapshot

ory · Mar 21, 2024 · c9d2dd3 · c9d2dd3
1 parent 6511eb1
commit c9d2dd3
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 1 deletion.
diff --git a/jsonnetsecure/jsonnet_test.go b/jsonnetsecure/jsonnet_test.go
@@ -60,7 +60,7 @@ func TestSecureVM(t *testing.T) {
 		t.Run("case=simple input", func(t *testing.T) {
 			// from https://jsonnet.org/learning/tutorial.html
 			snippet := `
-/* A C-style comment. */
+/* A C-style comment. \u0000 \000 \x00 */
 # A Python-style comment.
 {
   cocktails: {
@@ -100,6 +100,21 @@ func TestSecureVM(t *testing.T) {
 			})
 		})
 
+		t.Run("case=null bytes", func(t *testing.T) {
+			snippet := `{a: 0,`
+			snippet += "/*" + string(rune(0)) + "*/"
+			snippet += ` b: 1}`
+
+			assertEqualVMOutput(t, func(factory func(t *testing.T) VM) string {
+				vm := factory(t)
+				vm.ExtVar("nullvar", string(rune(0)))
+				vm.ExtCode("nullcode", string(rune(0)))
+				out, err := vm.EvaluateAnonymousSnippet("test", snippet)
+				assert.NoError(t, err)
+				return out
+			})
+		})
+
 		t.Run("case=ext variables", func(t *testing.T) {
 			assertEqualVMOutput(t, func(factory func(t *testing.T) VM) string {
 				vm := factory(t)

diff --git a/snapshotx/snapshot.go b/snapshotx/snapshot.go
@@ -4,6 +4,7 @@
 package snapshotx
 
 import (
+	"bytes"
 	"encoding/json"
 	"fmt"
 	"strings"
@@ -25,6 +26,7 @@ type (
 	}
 	exceptPaths      []string
 	exceptNestedKeys []string
+	replacement      struct{ str, replacement string }
 )
 
 func (e exceptPaths) apply(t *testing.T, raw []byte) []byte {
@@ -42,6 +44,10 @@ func (e exceptNestedKeys) apply(t *testing.T, raw []byte) []byte {
 	return deleteMatches(t, "", parsed, e, []string{}, raw)
 }
 
+func (r *replacement) apply(_ *testing.T, raw []byte) []byte {
+	return bytes.ReplaceAll(raw, []byte(r.str), []byte(r.replacement))
+}
+
 func ExceptPaths(keys ...string) ExceptOpt {
 	return exceptPaths(keys)
 }
@@ -50,6 +56,10 @@ func ExceptNestedKeys(nestedKeys ...string) ExceptOpt {
 	return exceptNestedKeys(nestedKeys)
 }
 
+func WithReplacement(str, replace string) ExceptOpt {
+	return &replacement{str: str, replacement: replace}
+}
+
 func SnapshotTJSON(t *testing.T, compare []byte, except ...ExceptOpt) {
 	t.Helper()
 	for _, e := range except {