Add yaml_utils.RemoveKey

Co-authored-by: Chris McDonnell <[email protected]>

Author: stefanhaller

pkg/utils/yaml_utils/yaml_utils.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
+	"slices"
 
 	"gopkg.in/yaml.v3"
 )
@@ -18,6 +19,19 @@ func lookupKey(node *yaml.Node, key string) (*yaml.Node, *yaml.Node) {
 	return nil, nil
 }
 
+// Returns the key and value if they were present
+func RemoveKey(node *yaml.Node, key string) (*yaml.Node, *yaml.Node) {
+	for i := 0; i < len(node.Content)-1; i += 2 {
+		if node.Content[i].Value == key {
+			key, value := node.Content[i], node.Content[i+1]
+			node.Content = slices.Delete(node.Content, i, i+2)
+			return key, value
+		}
+	}
+
+	return nil, nil
+}
+
 // Walks a yaml document from the root node to the specified path, and then applies the transformation to that node.
 // If the requested path is not defined in the document, no changes are made to the document.
 func TransformNode(rootNode *yaml.Node, path []string, transform func(node *yaml.Node) error) error {

pkg/utils/yaml_utils/yaml_utils_test.go

@@ -282,6 +282,55 @@ foo:
 	}
 }
 
+func TestRemoveKey(t *testing.T) {
+	tests := []struct {
+		name                 string
+		in                   string
+		key                  string
+		expectedOut          string
+		expectedRemovedKey   string
+		expectedRemovedValue string
+	}{
+		{
+			name: "Key not present",
+			in:   "foo: 1",
+			key:  "bar",
+		},
+		{
+			name:                 "Key present",
+			in:                   "foo: 1\nbar: 2\nbaz: 3\n",
+			key:                  "bar",
+			expectedOut:          "foo: 1\nbaz: 3\n",
+			expectedRemovedKey:   "bar",
+			expectedRemovedValue: "2",
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			node := unmarshalForTest(t, test.in)
+			removedKey, removedValue := RemoveKey(node.Content[0], test.key)
+			if test.expectedOut == "" {
+				unmodifiedOriginal := unmarshalForTest(t, test.in)
+				assert.Equal(t, unmodifiedOriginal, node)
+			} else {
+				result := marshalForTest(t, &node)
+				assert.Equal(t, test.expectedOut, result)
+			}
+			if test.expectedRemovedKey == "" {
+				assert.Nil(t, removedKey)
+			} else {
+				assert.Equal(t, test.expectedRemovedKey, removedKey.Value)
+			}
+			if test.expectedRemovedValue == "" {
+				assert.Nil(t, removedValue)
+			} else {
+				assert.Equal(t, test.expectedRemovedValue, removedValue.Value)
+			}
+		})
+	}
+}
+
 func unmarshalForTest(t *testing.T, input string) yaml.Node {
 	t.Helper()
 	var node yaml.Node