From 2161d095fa118b4c7efc7fff33475de4c6e2f45d Mon Sep 17 00:00:00 2001
From: Qirui <qiruiluo2022@u.northwestern.edu>
Date: Mon, 20 Mar 2023 11:42:38 +0800
Subject: [PATCH] Chore: add template restart policy tweak function

---
 pkg/apis/apps/validation/validation_test.go | 71 +++------------------
 1 file changed, 9 insertions(+), 62 deletions(-)

diff --git a/pkg/apis/apps/validation/validation_test.go b/pkg/apis/apps/validation/validation_test.go
index 6732757991c..8b606668ead 100644
--- a/pkg/apis/apps/validation/validation_test.go
+++ b/pkg/apis/apps/validation/validation_test.go
@@ -86,6 +86,12 @@ func tweakReplicas(replicas int32) statefulSetTweak {
 	}
 }
 
+func tweakTemplateRestartPolicy(rp api.RestartPolicy) statefulSetTweak {
+	return func(ss *apps.StatefulSet) {
+		ss.Spec.Template.Spec.RestartPolicy = rp
+	}
+}
+
 func TestValidateStatefulSet(t *testing.T) {
 	validLabels := map[string]string{"a": "b"}
 	validPodTemplate := api.PodTemplate{
@@ -354,80 +360,21 @@ func TestValidateStatefulSet(t *testing.T) {
 		},
 		{
 			name: "invalid restart policy 1",
-			set: apps.StatefulSet{
-				ObjectMeta: metav1.ObjectMeta{
-					Name:      "abc-123",
-					Namespace: metav1.NamespaceDefault,
-				},
-				Spec: apps.StatefulSetSpec{
-					PodManagementPolicy: apps.OrderedReadyPodManagement,
-					Selector:            &metav1.LabelSelector{MatchLabels: validLabels},
-					Template: api.PodTemplateSpec{
-						Spec: api.PodSpec{
-							RestartPolicy: api.RestartPolicyOnFailure,
-							DNSPolicy:     api.DNSClusterFirst,
-							Containers:    []api.Container{{Name: "ctr", Image: "image", ImagePullPolicy: "IfNotPresent"}},
-						},
-						ObjectMeta: metav1.ObjectMeta{
-							Labels: validLabels,
-						},
-					},
-					UpdateStrategy: apps.StatefulSetUpdateStrategy{Type: apps.RollingUpdateStatefulSetStrategyType},
-				},
-			},
+			set:  mkStatefulSet(&validPodTemplate, tweakTemplateRestartPolicy(api.RestartPolicyOnFailure)),
 			errs: field.ErrorList{
 				field.NotSupported(field.NewPath("spec", "template", "spec", "restartPolicy"), nil, nil),
 			},
 		},
 		{
 			name: "invalid restart policy 2",
-			set: apps.StatefulSet{
-				ObjectMeta: metav1.ObjectMeta{
-					Name:      "abc-123",
-					Namespace: metav1.NamespaceDefault,
-				},
-				Spec: apps.StatefulSetSpec{
-					PodManagementPolicy: apps.OrderedReadyPodManagement,
-					Selector:            &metav1.LabelSelector{MatchLabels: validLabels},
-					Template: api.PodTemplateSpec{
-						Spec: api.PodSpec{
-							RestartPolicy: api.RestartPolicyNever,
-							DNSPolicy:     api.DNSClusterFirst,
-							Containers:    []api.Container{{Name: "ctr", Image: "image", ImagePullPolicy: "IfNotPresent"}},
-						},
-						ObjectMeta: metav1.ObjectMeta{
-							Labels: validLabels,
-						},
-					},
-					UpdateStrategy: apps.StatefulSetUpdateStrategy{Type: apps.RollingUpdateStatefulSetStrategyType},
-				},
-			},
+			set:  mkStatefulSet(&validPodTemplate, tweakTemplateRestartPolicy(api.RestartPolicyNever)),
 			errs: field.ErrorList{
 				field.NotSupported(field.NewPath("spec", "template", "spec", "restartPolicy"), nil, nil),
 			},
 		},
 		{
 			name: "empty restart policy",
-			set: apps.StatefulSet{
-				ObjectMeta: metav1.ObjectMeta{
-					Name:      "abc-123",
-					Namespace: metav1.NamespaceDefault,
-				},
-				Spec: apps.StatefulSetSpec{
-					PodManagementPolicy: apps.OrderedReadyPodManagement,
-					Selector:            &metav1.LabelSelector{MatchLabels: validLabels},
-					Template: api.PodTemplateSpec{
-						Spec: api.PodSpec{
-							DNSPolicy:  api.DNSClusterFirst,
-							Containers: []api.Container{{Name: "ctr", Image: "image", ImagePullPolicy: "IfNotPresent"}},
-						},
-						ObjectMeta: metav1.ObjectMeta{
-							Labels: validLabels,
-						},
-					},
-					UpdateStrategy: apps.StatefulSetUpdateStrategy{Type: apps.RollingUpdateStatefulSetStrategyType},
-				},
-			},
+			set:  mkStatefulSet(&validPodTemplate, tweakTemplateRestartPolicy("")),
 			errs: field.ErrorList{
 				field.NotSupported(field.NewPath("spec", "template", "spec", "restartPolicy"), nil, nil),
 			},