Add UT

Author: pajakd

slice/internal/webhooks/statefulset_webhook_test.go

@@ -93,7 +93,7 @@ func TestStatefulSetDefault(t *testing.T) {
 						},
 						Spec: corev1.PodSpec{
 							NodeSelector: map[string]string{
-								"cloud.google.com/gke-tpu-accelerator": string(slice.TypeTpu7x),
+								core.TPUAcceleratorLabel: string(slice.TypeTpu7x),
 							},
 						},
 					},
@@ -111,8 +111,48 @@ func TestStatefulSetDefault(t *testing.T) {
 						},
 						Spec: corev1.PodSpec{
 							NodeSelector: map[string]string{
-								"cloud.google.com/gke-tpu-accelerator": string(slice.TypeTpu7x),
-								core.TPUTopologyAnnotation:             "4x4x12",
+								core.TPUAcceleratorLabel:   string(slice.TypeTpu7x),
+								core.TPUTopologyAnnotation: "4x4x12",
+							},
+						},
+					},
+				},
+			},
+		},
+		"relevant statefulset, with tpu topology node selector already present": {
+			sts: &appsv1.StatefulSet{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      baseName,
+					Namespace: baseNamespace,
+				},
+				Spec: appsv1.StatefulSetSpec{
+					Template: corev1.PodTemplateSpec{
+						ObjectMeta: metav1.ObjectMeta{
+							Annotations: map[string]string{core.TPUSliceTopologyAnnotation: "4x4x12"},
+						},
+						Spec: corev1.PodSpec{
+							NodeSelector: map[string]string{
+								core.TPUAcceleratorLabel:   string(slice.TypeTpu7x),
+								core.TPUTopologyAnnotation: "4x4x12",
+							},
+						},
+					},
+				},
+			},
+			wantSts: &appsv1.StatefulSet{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      baseName,
+					Namespace: baseNamespace,
+				},
+				Spec: appsv1.StatefulSetSpec{
+					Template: corev1.PodTemplateSpec{
+						ObjectMeta: metav1.ObjectMeta{
+							Annotations: map[string]string{core.TPUSliceTopologyAnnotation: "4x4x12"},
+						},
+						Spec: corev1.PodSpec{
+							NodeSelector: map[string]string{
+								core.TPUAcceleratorLabel:   string(slice.TypeTpu7x),
+								core.TPUTopologyAnnotation: "4x4x12",
 							},
 						},
 					},