diff --git a/pkg/utils/yaml_utils/yaml_utils.go b/pkg/utils/yaml_utils/yaml_utils.go new file mode 100644 index 000000000..9ed7ae875 --- /dev/null +++ b/pkg/utils/yaml_utils/yaml_utils.go @@ -0,0 +1,54 @@ +package yaml_utils + +import ( + "fmt" + + "gopkg.in/yaml.v3" +) + +// takes a yaml document in bytes, a path to a key, and a value to set. The value must be a scalar. +func UpdateYaml(yamlBytes []byte, path []string, value string) ([]byte, error) { + // Parse the YAML file. + var node yaml.Node + err := yaml.Unmarshal(yamlBytes, &node) + if err != nil { + return nil, fmt.Errorf("failed to parse YAML: %w", err) + } + + body := node.Content[0] + + updateYamlNode(body, path, value) + + // Convert the updated YAML node back to YAML bytes. + updatedYAMLBytes, err := yaml.Marshal(body) + if err != nil { + return nil, fmt.Errorf("failed to convert YAML node to bytes: %w", err) + } + + return updatedYAMLBytes, nil +} + +// Recursive function to update the YAML node. +func updateYamlNode(node *yaml.Node, path []string, value string) { + if len(path) == 0 { + node.Value = value + return + } + + key := path[0] + for i := 0; i < len(node.Content)-1; i += 2 { + if node.Content[i].Value == key { + updateYamlNode(node.Content[i+1], path[1:], value) + return + } + } + + // if the key doesn't exist, we'll add it + node.Content = append(node.Content, &yaml.Node{ + Kind: yaml.ScalarNode, + Value: key, + }, &yaml.Node{ + Kind: yaml.ScalarNode, + Value: value, + }) +} diff --git a/pkg/utils/yaml_utils/yaml_utils_test.go b/pkg/utils/yaml_utils/yaml_utils_test.go new file mode 100644 index 000000000..4bdfeb432 --- /dev/null +++ b/pkg/utils/yaml_utils/yaml_utils_test.go @@ -0,0 +1,64 @@ +package yaml_utils + +import "testing" + +func TestUpdateYaml(t *testing.T) { + tests := []struct { + name string + in string + path []string + value string + expectedOut string + expectedErr string + }{ + { + name: "update value", + in: "foo: bar\n", + path: []string{"foo"}, + value: "baz", + expectedOut: "foo: baz\n", + expectedErr: "", + }, + { + name: "add new key and value", + in: "foo: bar\n", + path: []string{"foo2"}, + value: "baz", + expectedOut: "foo: bar\nfoo2: baz\n", + expectedErr: "", + }, + { + name: "preserve inline comment", + in: "foo: bar # my comment\n", + path: []string{"foo2"}, + value: "baz", + expectedOut: "foo: bar # my comment\nfoo2: baz\n", + expectedErr: "", + }, + { + name: "nested update", + in: "foo:\n bar: baz\n", + path: []string{"foo", "bar"}, + value: "qux", + // indentation is not preserved. See https://github.com/go-yaml/yaml/issues/899 + expectedOut: "foo:\n bar: qux\n", + expectedErr: "", + }, + } + + for _, test := range tests { + test := test + t.Run(test.name, func(t *testing.T) { + out, err := UpdateYaml([]byte(test.in), test.path, test.value) + if test.expectedErr != "" { + if err == nil { + t.Errorf("expected error %q but got none", test.expectedErr) + } + } else if err != nil { + t.Errorf("unexpected error: %v", err) + } else if string(out) != test.expectedOut { + t.Errorf("expected %q but got %q", test.expectedOut, string(out)) + } + }) + } +}