Add KubeObject setters that keep formatting to "kubeobject" lib

.gitignore

@@ -17,3 +17,4 @@
 # Tox environments
 .tox/
 *.dic
+/.vscode

go.mod

@@ -11,6 +11,7 @@ require (
 	github.com/stretchr/testify v1.8.1
 	k8s.io/api v0.26.1
 	k8s.io/apimachinery v0.26.3
+	sigs.k8s.io/kustomize/kyaml v0.13.9
 	sigs.k8s.io/yaml v1.3.0
 )
 
@@ -48,6 +49,5 @@ require (
 	sigs.k8s.io/controller-runtime v0.14.5 // indirect
 	sigs.k8s.io/json v0.0.0-20220713155537-f223a00ba0e2 // indirect
 	sigs.k8s.io/kustomize/api v0.12.1 // indirect
-	sigs.k8s.io/kustomize/kyaml v0.13.9 // indirect
 	sigs.k8s.io/structured-merge-diff/v4 v4.2.3 // indirect
 )

krm-functions/lib/kubeobject/kubeobject.go

@@ -18,8 +18,11 @@ package kubeobject
 
 import (
 	"fmt"
+	"reflect"
+	"unsafe"
 
 	"github.com/GoogleContainerTools/kpt-functions-sdk/go/fn"
+	"sigs.k8s.io/kustomize/kyaml/yaml"
 )
 
 type KubeObjectExt[T1 any] struct {
@@ -63,3 +66,136 @@ func NewFromGoStruct[T1 any](x any) (*KubeObjectExt[T1], error) {
 	}
 	return NewFromKubeObject[T1](o)
 }
+
+// SetSpec sets the `spec` field of a KubeObjectExt to the value of `newSpec`,
+// while trying to keep as much formatting as possible
+func (o *KubeObjectExt[T1]) SetSpec(newSpec interface{}) error {
+	return setNestedFieldKeepFormatting(&o.KubeObject.SubObject, newSpec, "spec")
+}
+
+// SetStatus sets the `status` field of a KubeObjectExt to the value of `newStatus`,
+// while trying to keep as much formatting as possible
+func (o *KubeObjectExt[T1]) SetStatus(newStatus interface{}) error {
+	return setNestedFieldKeepFormatting(&o.KubeObject.SubObject, newStatus, "status")
+}
+
+// NOTE: the following functions are considered as "methods" of KubeObject,
+// and thus nill checking of `obj` was omitted intentionally:
+// the caller is responsible for ensuring that `obj` is not nil`
+
+// setNestedFieldKeepFormatting is similar to KubeObject.SetNestedField(), but keeps the
+// comments and the order of fields in the YAML wherever it is possible.
+//
+// NOTE: This functionality should be solved in the upstream SDK.
+// Merging the code below to the upstream SDK is in progress and tracked in this issue:
+// https://github.com/GoogleContainerTools/kpt/issues/3923
+func setNestedFieldKeepFormatting(obj *fn.SubObject, value interface{}, field string) error {
+	oldNode := yamlNodeOf(obj.UpsertMap(field))
+	err := obj.SetNestedField(value, field)
+	if err != nil {
+		return err
+	}
+	newNode := yamlNodeOf(obj.GetMap(field))
+
+	restoreFieldOrder(oldNode, newNode)
+	deepCopyComments(oldNode, newNode)
+	return nil
+}
+
+///////////////// internals
+
+func shallowCopyComments(src, dst *yaml.Node) {
+	dst.HeadComment = src.HeadComment
+	dst.LineComment = src.LineComment
+	dst.FootComment = src.FootComment
+}
+
+func deepCopyComments(src, dst *yaml.Node) {
+	if src.Kind != dst.Kind {
+		return
+	}
+	shallowCopyComments(src, dst)
+	if dst.Kind == yaml.MappingNode {
+		if (len(src.Content)%2 != 0) || (len(dst.Content)%2 != 0) {
+			panic("unexpected number of children for YAML map")
+		}
+		for i := 0; i < len(dst.Content); i += 2 {
+			dstKeyNode := dst.Content[i]
+			key, ok := asString(dstKeyNode)
+			if !ok {
+				continue
+			}
+
+			j, ok := findKey(src, key)
+			if !ok {
+				continue
+			}
+			srcKeyNode, srcValueNode := src.Content[j], src.Content[j+1]
+			dstValueNode := dst.Content[i+1]
+			shallowCopyComments(srcKeyNode, dstKeyNode)
+			deepCopyComments(srcValueNode, dstValueNode)
+		}
+	}
+}
+
+func restoreFieldOrder(src, dst *yaml.Node) {
+	if (src.Kind != dst.Kind) || (dst.Kind != yaml.MappingNode) {
+		return
+	}
+	if (len(src.Content)%2 != 0) || (len(dst.Content)%2 != 0) {
+		panic("unexpected number of children for YAML map")
+	}
+
+	nextInDst := 0
+	for i := 0; i < len(src.Content); i += 2 {
+		key, ok := asString(src.Content[i])
+		if !ok {
+			continue
+		}
+
+		j, ok := findKey(dst, key)
+		if !ok {
+			continue
+		}
+		if j != nextInDst {
+			dst.Content[j], dst.Content[nextInDst] = dst.Content[nextInDst], dst.Content[j]
+			dst.Content[j+1], dst.Content[nextInDst+1] = dst.Content[nextInDst+1], dst.Content[j+1]
+		}
+		nextInDst += 2
+
+		srcValueNode := src.Content[i+1]
+		dstValueNode := dst.Content[nextInDst-1]
+		restoreFieldOrder(srcValueNode, dstValueNode)
+	}
+}
+
+func asString(node *yaml.Node) (string, bool) {
+	if node.Kind == yaml.ScalarNode && (node.Tag == "!!str" || node.Tag == "") {
+		return node.Value, true
+	}
+	return "", false
+}
+
+func findKey(m *yaml.Node, key string) (int, bool) {
+	children := m.Content
+	if len(children)%2 != 0 {
+		panic("unexpected number of children for YAML map")
+	}
+	for i := 0; i < len(children); i += 2 {
+		keyNode := children[i]
+		k, ok := asString(keyNode)
+		if ok && k == key {
+			return i, true
+		}
+	}
+	return 0, false
+}
+
+// This is a temporary workaround until SetNestedFieldKeppFormatting functionality is merged into the upstream SDK
+// The merge process has already started and tracked in this issue: https://github.com/GoogleContainerTools/kpt/issues/3923
+func yamlNodeOf(obj *fn.SubObject) *yaml.Node {
+	internalObj := reflect.ValueOf(*obj).FieldByName("obj")
+	nodePtr := internalObj.Elem().FieldByName("node")
+	nodePtr = reflect.NewAt(nodePtr.Type(), unsafe.Pointer(nodePtr.UnsafeAddr())).Elem()
+	return nodePtr.Interface().(*yaml.Node)
+}

krm-functions/lib/kubeobject/kubeobject_test.go

@@ -17,10 +17,15 @@ limitations under the License.
 package kubeobject
 
 import (
+	"os"
+	"strings"
 	"testing"
 
 	"github.com/GoogleContainerTools/kpt-functions-sdk/go/fn"
-	v1 "k8s.io/api/apps/v1"
+	"github.com/GoogleContainerTools/kpt-functions-sdk/go/fn/testhelpers"
+	testlib "github.com/nephio-project/nephio/krm-functions/lib/test"
+	appsv1 "k8s.io/api/apps/v1"
+	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"sigs.k8s.io/yaml"
 )
@@ -93,7 +98,7 @@ func TestNewFromKubeObject(t *testing.T) {
 		},
 	}
 	for _, tt := range testItems {
-		deploymentReceived := v1.Deployment{
+		deploymentReceived := appsv1.Deployment{
 			TypeMeta: metav1.TypeMeta{
 				APIVersion: tt.input.gv,
 				Kind:       tt.input.kind,
@@ -102,7 +107,7 @@ func TestNewFromKubeObject(t *testing.T) {
 				Name:      tt.input.name,
 				Namespace: tt.input.namespace,
 			},
-			Spec: v1.DeploymentSpec{
+			Spec: appsv1.DeploymentSpec{
 				Replicas: &tt.input.replicas,
 				Paused:   tt.input.paused,
 				Selector: &metav1.LabelSelector{
@@ -115,7 +120,7 @@ func TestNewFromKubeObject(t *testing.T) {
 			t.Errorf("YAML Marshal error: %s", err)
 		}
 		deploymentKubeObject, _ := fn.ParseKubeObject(b)
-		deploymentKubeObjectParser, _ := NewFromKubeObject[v1.Deployment](deploymentKubeObject)
+		deploymentKubeObjectParser, _ := NewFromKubeObject[appsv1.Deployment](deploymentKubeObject)
 		if deploymentKubeObjectParser.SubObject != deploymentKubeObject.SubObject {
 			t.Errorf("-want%s, +got:\n%s", deploymentKubeObjectParser.String(), deploymentKubeObject.String())
 		}
@@ -198,7 +203,7 @@ func TestNewFromYaml(t *testing.T) {
 		},
 	}
 	for _, tt := range testItems {
-		deploymentReceived := v1.Deployment{
+		deploymentReceived := appsv1.Deployment{
 			TypeMeta: metav1.TypeMeta{
 				APIVersion: tt.input.gv,
 				Kind:       tt.input.kind,
@@ -207,7 +212,7 @@ func TestNewFromYaml(t *testing.T) {
 				Name:      tt.input.name,
 				Namespace: tt.input.namespace,
 			},
-			Spec: v1.DeploymentSpec{
+			Spec: appsv1.DeploymentSpec{
 				Replicas: &tt.input.replicas,
 				Paused:   tt.input.paused,
 				Selector: &metav1.LabelSelector{
@@ -219,7 +224,7 @@ func TestNewFromYaml(t *testing.T) {
 		if err != nil {
 			t.Errorf("YAML Marshal error: %s", err)
 		}
-		deploymentKubeObjectParser, _ := NewFromYaml[v1.Deployment](b)
+		deploymentKubeObjectParser, _ := NewFromYaml[appsv1.Deployment](b)
 
 		if deploymentKubeObjectParser.String() != string(b) {
 			t.Errorf("-want%s, +got:\n%s", string(b), deploymentKubeObjectParser.String())
@@ -303,7 +308,7 @@ func TestNewFromGoStruct(t *testing.T) {
 		},
 	}
 	for _, tt := range testItems {
-		deploymentReceived := v1.Deployment{
+		deploymentReceived := appsv1.Deployment{
 			TypeMeta: metav1.TypeMeta{
 				APIVersion: tt.input.gv,
 				Kind:       tt.input.kind,
@@ -312,15 +317,15 @@ func TestNewFromGoStruct(t *testing.T) {
 				Name:      tt.input.name,
 				Namespace: tt.input.namespace,
 			},
-			Spec: v1.DeploymentSpec{
+			Spec: appsv1.DeploymentSpec{
 				Replicas: &tt.input.replicas,
 				Paused:   tt.input.paused,
 				Selector: &metav1.LabelSelector{
 					MatchLabels: tt.input.selector,
 				},
 			},
 		}
-		deploymentKubeObject, _ := NewFromGoStruct[v1.Deployment](deploymentReceived)
+		deploymentKubeObject, _ := NewFromGoStruct[appsv1.Deployment](deploymentReceived)
 
 		s, _, err := deploymentKubeObject.NestedString([]string{"metadata", "name"}...)
 		if err != nil {
@@ -332,8 +337,105 @@ func TestNewFromGoStruct(t *testing.T) {
 	}
 
 	// test with nil input
-	_, err := NewFromGoStruct[v1.Deployment](nil)
+	_, err := NewFromGoStruct[appsv1.Deployment](nil)
 	if err == nil {
 		t.Errorf("NewFromGoStruct(nil) doesn't return with an error")
 	}
 }
+
+func compareKubeObjectWithExpectedYaml(t *testing.T, obj *fn.KubeObject, inputFile string) {
+	actualYAML := strings.TrimSpace(obj.String())
+	expectedFile := testlib.InsertBeforeExtension(inputFile, "_expected")
+	expectedYAML := strings.TrimSpace(string(testhelpers.MustReadFile(t, expectedFile)))
+
+	if actualYAML != expectedYAML {
+		t.Errorf(`mismatch in expected and actual KubeObject YAML:
+--- want: -----
+%v
+--- got: ----
+%v
+----------------`, expectedYAML, actualYAML)
+		os.WriteFile(testlib.InsertBeforeExtension(inputFile, "_actual"), []byte(actualYAML), 0666)
+	}
+
+}
+
+func TestSetNestedFieldKeepFormatting(t *testing.T) {
+	testfiles := []string{"testdata/comments.yaml"}
+	for _, inputFile := range testfiles {
+		t.Run(inputFile, func(t *testing.T) {
+			obj := testlib.MustParseKubeObject(t, inputFile)
+
+			var deploy appsv1.Deployment
+			if err := obj.As(&deploy); err != nil {
+				t.Fatalf("couldn't convert object to Deployment: %v", err)
+			}
+
+			deploy.Spec.Replicas = nil                                              // delete Replicas field if present
+			deploy.Spec.Template.Spec.RestartPolicy = corev1.RestartPolicyOnFailure // update field value
+
+			err := setNestedFieldKeepFormatting(&obj.SubObject, deploy.Spec, "spec")
+			if err != nil {
+				t.Errorf("unexpected error in SetNestedFieldKeepFormatting: %v", err)
+			}
+
+			compareKubeObjectWithExpectedYaml(t, obj, inputFile)
+		})
+	}
+}
+
+func TestKubeObjectExtSetSpec(t *testing.T) {
+	testfiles := []string{"testdata/comments.yaml"}
+	for _, inputFile := range testfiles {
+		t.Run(inputFile, func(t *testing.T) {
+			obj := testlib.MustParseKubeObject(t, inputFile)
+
+			koe, err := NewFromKubeObject[appsv1.Deployment](obj)
+			if err != nil {
+				t.Errorf("unexpected error: %v", err)
+			}
+			deploy, err := koe.GetGoStruct()
+			if err != nil {
+				t.Errorf("unexpected error: %v", err)
+			}
+
+			deploy.Spec.Replicas = nil                                              // delete Replicas field if present
+			deploy.Spec.Template.Spec.RestartPolicy = corev1.RestartPolicyOnFailure // update field value
+			err = koe.SetSpec(&deploy.Spec)
+			if err != nil {
+				t.Errorf("unexpected error in SetSpec: %v", err)
+			}
+
+			compareKubeObjectWithExpectedYaml(t, obj, inputFile)
+		})
+	}
+}
+
+func TestKubeObjectExtSetStatus(t *testing.T) {
+	testfiles := []string{
+		"testdata/status_comments.yaml",
+		"testdata/empty_status.yaml",
+	}
+	for _, inputFile := range testfiles {
+		t.Run(inputFile, func(t *testing.T) {
+			obj := testlib.MustParseKubeObject(t, inputFile)
+
+			koe, err := NewFromKubeObject[appsv1.Deployment](obj)
+			if err != nil {
+				t.Errorf("unexpected error: %v", err)
+			}
+			deploy, err := koe.GetGoStruct()
+			if err != nil {
+				t.Errorf("unexpected error: %v", err)
+			}
+
+			deploy.Status.AvailableReplicas = 0
+			err = koe.SetStatus(deploy.Status)
+			if err != nil {
+				t.Errorf("unexpected error in SetStatus: %v", err)
+			}
+
+			compareKubeObjectWithExpectedYaml(t, obj, inputFile)
+		})
+	}
+}

krm-functions/lib/kubeobject/testdata/.gitignore

@@ -0,0 +1 @@
+/*_actual.yaml