Fix comments handling in addOrReplaceValue function (#1039)

pkg/mutation/mutate.go

@@ -182,6 +182,9 @@ func addOrReplaceValue(node *yaml.Node, splits []string, value *yaml.Node) error
 	}
 	for _, node := range nodes {
 		if node.Kind == yaml.ScalarNode {
+			if value.LineComment == "" {
+				value.LineComment = node.LineComment // keep the original comment if override is not provided
+			}
 			// Overwrite an existing scalar value with a new value (whatever kind).
 			*node = *value
 		} else if node.Kind == yaml.MappingNode && value.Kind == yaml.MappingNode {

pkg/mutation/mutate_test.go

@@ -107,31 +107,26 @@ obj:
 			},
 			Path: "/obj/foo",
 		},
-		mutated: `obj:
+		mutated: `
+obj:
   foo:
     bar:
       - c
       - d
   baz: quux
 `,
 	}, {
-		original: `
-foo: bar
-`,
+		original: `foo: bar`,
 		patch: config.Mutation{
 			Op:      "replace",
 			Value:   "baz",
 			Path:    "/foo",
 			Comment: "# We set this to baz",
 		},
-		mutated: `
-foo: baz # We set this to baz
-`,
+		mutated: `foo: baz # We set this to baz`,
 		message: "Expected a comment to appear",
 	}, {
-		original: `
-foo: bar
-`,
+		original: `foo: bar`,
 		patch: config.Mutation{
 			Op: "add",
 			Value: map[string]interface{}{
@@ -167,6 +162,25 @@ extra:
   baz: quux
 `,
 		message: "Expected a comment to appear next to an object",
+	}, {
+		original: `foo: bar # we should keep this comment`,
+		patch: config.Mutation{
+			Op:    "replace",
+			Value: "baz",
+			Path:  "/foo",
+		},
+		mutated: `foo: baz # we should keep this comment`,
+		message: "Expected a comment to be kept",
+	}, {
+		original: `foo: bar # we should override this comment`,
+		patch: config.Mutation{
+			Op:      "replace",
+			Value:   "baz",
+			Path:    "/foo",
+			Comment: "override",
+		},
+		mutated: `foo: baz # override`,
+		message: "Expected a comment to overridden",
 	},
 }