From 39a86844f24169b632513ef188d90f1ba6b72665 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=8E=8B=E6=AF=93?= <wangyu.luck@bytedance.com>
Date: Thu, 18 Jan 2024 20:57:41 +0800
Subject: [PATCH] feat: retain pod spec containers resources when sync pod
 reource

---
 pkg/controllers/sync/dispatch/retain.go      | 12 ++++
 pkg/controllers/sync/dispatch/retain_test.go | 69 ++++++++++++++++++++
 2 files changed, 81 insertions(+)

diff --git a/pkg/controllers/sync/dispatch/retain.go b/pkg/controllers/sync/dispatch/retain.go
index 78a722b9..e2c833ee 100644
--- a/pkg/controllers/sync/dispatch/retain.go
+++ b/pkg/controllers/sync/dispatch/retain.go
@@ -470,6 +470,18 @@ func retainContainers(desiredContainers, clusterContainers []interface{}) error
 }
 
 func retainContainer(desiredContainer, clusterContainer map[string]interface{}) error {
+	resources, exists, err := unstructured.NestedFieldNoCopy(clusterContainer, "resources")
+	if err != nil {
+		return err
+	}
+	if exists {
+		if err = unstructured.SetNestedField(desiredContainer, resources, "resources"); err != nil {
+			return err
+		}
+	} else {
+		unstructured.RemoveNestedField(desiredContainer, "resources")
+	}
+
 	if _, _, exists := findServiceAccountVolumeMount(desiredContainer); !exists {
 		if volumeMnt, idx, exists := findServiceAccountVolumeMount(clusterContainer); exists {
 			// The logic for retaining service account volume mounts is the same as retaining service account volumes.
diff --git a/pkg/controllers/sync/dispatch/retain_test.go b/pkg/controllers/sync/dispatch/retain_test.go
index 6e9bbf51..328e0928 100644
--- a/pkg/controllers/sync/dispatch/retain_test.go
+++ b/pkg/controllers/sync/dispatch/retain_test.go
@@ -21,6 +21,7 @@ are Copyright 2023 The KubeAdmiral Authors.
 package dispatch
 
 import (
+	"reflect"
 	"testing"
 
 	"github.com/onsi/gomega"
@@ -201,3 +202,71 @@ func TestMergeStringMaps(t *testing.T) {
 		})
 	}
 }
+
+func Test_retainContainer(t *testing.T) {
+	type args struct {
+		desiredContainer map[string]interface{}
+		clusterContainer map[string]interface{}
+	}
+	tests := []struct {
+		name    string
+		args    args
+		wantErr bool
+	}{
+		{
+			name: "nil retain test",
+			args: args{
+				desiredContainer: map[string]interface{}{
+					"name": "container-1",
+					"resources": map[string]interface{}{
+						"cpu":    "500m",
+						"memory": "512Mi",
+					}},
+				clusterContainer: map[string]interface{}{
+					"name": "container-1",
+				},
+			},
+		},
+		{
+			name: "empty retain test",
+			args: args{
+				desiredContainer: map[string]interface{}{
+					"name": "container-1",
+					"resources": map[string]interface{}{
+						"cpu":    "500m",
+						"memory": "512Mi",
+					}},
+				clusterContainer: map[string]interface{}{
+					"name":      "container-1",
+					"resources": map[string]interface{}{}},
+			},
+		},
+		{
+			name: "normal retain test",
+			args: args{
+				desiredContainer: map[string]interface{}{
+					"name": "container-1",
+					"resources": map[string]interface{}{
+						"cpu":    "500m",
+						"memory": "512Mi",
+					}},
+				clusterContainer: map[string]interface{}{
+					"name": "container-1",
+					"resources": map[string]interface{}{
+						"cpu":    "100m",
+						"memory": "100Mi",
+					}},
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if err := retainContainer(tt.args.desiredContainer, tt.args.clusterContainer); (err != nil) != tt.wantErr {
+				t.Errorf("retainContainer() error = %v, wantErr %v", err, tt.wantErr)
+			}
+			if !reflect.DeepEqual(tt.args.desiredContainer, tt.args.clusterContainer) {
+				t.Errorf("retainContainer did not retain the resources field correctly")
+			}
+		})
+	}
+}