From 34b679e0cc49ad70e9d8e38af07ec62c86eff494 Mon Sep 17 00:00:00 2001
From: Zhanwei Li <zhanweelee@gmail.com>
Date: Thu, 20 Feb 2025 14:23:49 +0800
Subject: [PATCH] feat: Add mustToYaml and mustToJson template functions

Introduces two new template functions that marshal data to YAML and
JSON, respectively, and panic on errors. This allows for strict
validation of template output formats.

Signed-off-by: Zhanwei Li <zhanweelee@gmail.com>
---
 pkg/engine/funcs.go      | 28 ++++++++++++++++++++++++++++
 pkg/engine/funcs_test.go | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+)

diff --git a/pkg/engine/funcs.go b/pkg/engine/funcs.go
index d03a818c2..c1f590018 100644
--- a/pkg/engine/funcs.go
+++ b/pkg/engine/funcs.go
@@ -51,10 +51,12 @@ func funcMap() template.FuncMap {
 		"toToml":        toTOML,
 		"fromToml":      fromTOML,
 		"toYaml":        toYAML,
+		"mustToYaml":    mustToYAML,
 		"toYamlPretty":  toYAMLPretty,
 		"fromYaml":      fromYAML,
 		"fromYamlArray": fromYAMLArray,
 		"toJson":        toJSON,
+		"mustToJson":    mustToJSON,
 		"fromJson":      fromJSON,
 		"fromJsonArray": fromJSONArray,
 
@@ -91,6 +93,19 @@ func toYAML(v interface{}) string {
 	return strings.TrimSuffix(string(data), "\n")
 }
 
+// mustToYAML takes an interface, marshals it to yaml, and returns a string.
+// It will panic if there is an error.
+//
+// This is designed to be called from a template when need to ensure that the
+// output YAML is valid.
+func mustToYAML(v interface{}) string {
+	data, err := yaml.Marshal(v)
+	if err != nil {
+		panic(err)
+	}
+	return strings.TrimSuffix(string(data), "\n")
+}
+
 func toYAMLPretty(v interface{}) string {
 	var data bytes.Buffer
 	encoder := goYaml.NewEncoder(&data)
@@ -176,6 +191,19 @@ func toJSON(v interface{}) string {
 	return string(data)
 }
 
+// mustToJSON takes an interface, marshals it to json, and returns a string.
+// It will panic if there is an error.
+//
+// This is designed to be called from a template when need to ensure that the
+// output JSON is valid.
+func mustToJSON(v interface{}) string {
+	data, err := json.Marshal(v)
+	if err != nil {
+		panic(err)
+	}
+	return string(data)
+}
+
 // fromJSON converts a JSON document into a map[string]interface{}.
 //
 // This is not a general-purpose JSON parser, and will not parse all valid
diff --git a/pkg/engine/funcs_test.go b/pkg/engine/funcs_test.go
index a4f4d604f..99edf5ae9 100644
--- a/pkg/engine/funcs_test.go
+++ b/pkg/engine/funcs_test.go
@@ -135,6 +135,43 @@ keyInElement1 = "valueInElement1"`,
 		assert.NoError(t, err)
 		assert.Equal(t, tt.expect, b.String(), tt.tpl)
 	}
+
+	loopMap := map[string]interface{}{
+		"foo": "bar",
+	}
+	loopMap["loop"] = []interface{}{loopMap}
+
+	mustFuncsTests := []struct {
+		tpl    string
+		expect interface{}
+		vars   interface{}
+	}{{
+		tpl:  `{{ mustToYaml . }}`,
+		vars: loopMap,
+	}, {
+		tpl:  `{{ mustToJson . }}`,
+		vars: loopMap,
+	}, {
+		tpl:    `{{ toYaml . }}`,
+		expect: "", // should return empty string and swallow error
+		vars:   loopMap,
+	}, {
+		tpl:    `{{ toJson . }}`,
+		expect: "", // should return empty string and swallow error
+		vars:   loopMap,
+	},
+	}
+
+	for _, tt := range mustFuncsTests {
+		var b strings.Builder
+		err := template.Must(template.New("test").Funcs(funcMap()).Parse(tt.tpl)).Execute(&b, tt.vars)
+		if tt.expect != nil {
+			assert.NoError(t, err)
+			assert.Equal(t, tt.expect, b.String(), tt.tpl)
+		} else {
+			assert.Error(t, err)
+		}
+	}
 }
 
 // This test to check a function provided by sprig is due to a change in a