From a81bd296742cadef17b9040b579e10b049619d0b Mon Sep 17 00:00:00 2001
From: Vitor Rodrigo Vezani <vitor.vezani@fairwinds.com>
Date: Fri, 15 Mar 2024 11:43:35 -0300
Subject: [PATCH] Fix comments handling in addOrReplaceValue function (#1039)

---
 pkg/mutation/mutate.go      |  3 +++
 pkg/mutation/mutate_test.go | 34 ++++++++++++++++++++++++----------
 2 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/pkg/mutation/mutate.go b/pkg/mutation/mutate.go
index 586a4af7d..30b8d4985 100644
--- a/pkg/mutation/mutate.go
+++ b/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 {
diff --git a/pkg/mutation/mutate_test.go b/pkg/mutation/mutate_test.go
index 921b73a3e..a7b81c007 100644
--- a/pkg/mutation/mutate_test.go
+++ b/pkg/mutation/mutate_test.go
@@ -107,7 +107,8 @@ obj:
 			},
 			Path: "/obj/foo",
 		},
-		mutated: `obj:
+		mutated: `
+obj:
   foo:
     bar:
       - c
@@ -115,23 +116,17 @@ obj:
   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",
 	},
 }