feat(coll): New coll.Set and coll.Unset functions (#2118)

Signed-off-by: Dave Henderson <dhenderson@gmail.com>

3 files changed, 58 insertions, 0 deletions

diff --git a/internal/funcs/coll.go b/internal/funcs/coll.go
index 945f77d3..0f87239d 100644
--- a/internal/funcs/coll.go
+++ b/internal/funcs/coll.go
@@ -33,6 +33,8 @@ func CreateCollFuncs(ctx context.Context) map[string]interface{} {
 	f["jsonpath"] = ns.JSONPath
 	f["jq"] = ns.JQ
 	f["flatten"] = ns.Flatten
+	f["set"] = ns.Set
+	f["unset"] = ns.Unset
 	return f
 }
 
@@ -221,3 +223,17 @@ func (CollFuncs) Omit(args ...interface{}) (map[string]interface{}, error) {
 	}
 	return coll.Omit(m, keys...), nil
 }
+
+// Set -
+func (CollFuncs) Set(key string, value interface{}, m map[string]interface{}) (map[string]interface{}, error) {
+	m[key] = value
+
+	return m, nil
+}
+
+// Unset -
+func (CollFuncs) Unset(key string, m map[string]interface{}) (map[string]interface{}, error) {
+	delete(m, key)
+
+	return m, nil
+}
diff --git a/internal/funcs/coll_test.go b/internal/funcs/coll_test.go
index fe0278d7..b075d078 100644
--- a/internal/funcs/coll_test.go
+++ b/internal/funcs/coll_test.go
@@ -216,3 +216,35 @@ func TestGoSlice(t *testing.T) {
 	assert.Equal(t, reflect.TypeOf([]string{}), out.Type())
 	assert.EqualValues(t, []string{"bar", "baz"}, out.Interface())
 }
+
+func TestCollFuncs_Set(t *testing.T) {
+	t.Parallel()
+
+	c := &CollFuncs{}
+
+	m := map[string]interface{}{"foo": "bar"}
+	out, err := c.Set("foo", "baz", m)
+	require.NoError(t, err)
+	assert.EqualValues(t, map[string]interface{}{"foo": "baz"}, out)
+
+	// m was modified so foo is now baz
+	out, err = c.Set("bar", "baz", m)
+	require.NoError(t, err)
+	assert.EqualValues(t, map[string]interface{}{"foo": "baz", "bar": "baz"}, out)
+}
+
+func TestCollFuncs_Unset(t *testing.T) {
+	t.Parallel()
+
+	c := &CollFuncs{}
+
+	m := map[string]interface{}{"foo": "bar"}
+	out, err := c.Unset("foo", m)
+	require.NoError(t, err)
+	assert.Empty(t, out)
+
+	// no-op
+	out, err = c.Unset("bar", m)
+	require.NoError(t, err)
+	assert.Empty(t, out)
+}
diff --git a/internal/tests/integration/collection_test.go b/internal/tests/integration/collection_test.go
index 45c33ad1..7ea15c1e 100644
--- a/internal/tests/integration/collection_test.go
+++ b/internal/tests/integration/collection_test.go
@@ -105,3 +105,13 @@ func TestColl_JQ(t *testing.T) {
 	inOutTest(t, `{{ coll.JQ ".foo" (dict "foo" 1 "bar" 2 "baz" 3) }}`, "1")
 	inOutTest(t, `{{ coll.Slice "one" 2 "three" 4.0 | jq ".[2]" }}`, `three`)
 }
+
+func TestColl_Set(t *testing.T) {
+	inOutTest(t, `{{ $dict := dict "foo" 1 }}{{ coll.Set "bar" 2 $dict }}`, "map[bar:2 foo:1]")
+	inOutTest(t, `{{ dict "foo" 1 | coll.Set "foo" 2 }}`, "map[foo:2]")
+}
+
+func TestColl_Unset(t *testing.T) {
+	inOutTest(t, `{{ $dict := dict "foo" 1 "bar" 2 }}{{ coll.Unset "bar" $dict }}`, "map[foo:1]")
+	inOutTest(t, `{{ dict "foo" 1 "bar" 2 | coll.Unset "foo" }}`, "map[bar:2]")
+}