extended e2e test to validate from UniquePerReplica to Shared

kubernetes-sigs · Sep 4, 2024 · e4d4b04 · e4d4b04
1 parent 5c45652
commit e4d4b04
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 6 deletions.
diff --git a/pkg/controllers/leaderworkerset_controller.go b/pkg/controllers/leaderworkerset_controller.go
@@ -106,7 +106,7 @@ func (r *LeaderWorkerSetReconciler) Reconcile(ctx context.Context, req ctrl.Requ
 	}
 
 	// Create headless service if it does not exist.
-	if err := r.createMultipleHeadlessServices(ctx, lws, replicas); err != nil {
+	if err := r.reconcileHeadlessServices(ctx, lws, replicas); err != nil {
 		log.Error(err, "Creating headless service.")
 		r.Record.Eventf(lws, corev1.EventTypeWarning, FailedCreate,
 			fmt.Sprintf("Failed to create headless service for error: %v", err))
@@ -122,8 +122,9 @@ func (r *LeaderWorkerSetReconciler) Reconcile(ctx context.Context, req ctrl.Requ
 	return ctrl.Result{}, nil
 }
 
-func (r *LeaderWorkerSetReconciler) createMultipleHeadlessServices(ctx context.Context, lws *leaderworkerset.LeaderWorkerSet, replicas int32) error {
-	if lws.Spec.NetworkConfig.SubdomainPolicy == leaderworkerset.SubdomainShared || lws.Spec.NetworkConfig == nil {
+func (r *LeaderWorkerSetReconciler) reconcileHeadlessServices(ctx context.Context, lws *leaderworkerset.LeaderWorkerSet, replicas int32) error {
+
+	if lws.Spec.NetworkConfig == nil || lws.Spec.NetworkConfig.SubdomainPolicy == leaderworkerset.SubdomainShared {
 		if err := r.createHeadlessServiceIfNotExists(ctx, lws, lws.Name, map[string]string{leaderworkerset.SetNameLabelKey: lws.Name}); err != nil {
 			return err
 		}
@@ -203,6 +204,7 @@ func (r *LeaderWorkerSetReconciler) deleteMultipleHeadlessSevices(ctx context.Co
 	if len(leaderPodList.Items) == 0 {
 		return nil
 	}
+
 	subdomainPolicy, foundSubdomainPolicy := leaderPodList.Items[0].Annotations[leaderworkerset.SubdomainPolicyAnnotationKey]
 	if foundSubdomainPolicy && subdomainPolicy == string(leaderworkerset.SubdomainUniquePerReplica) {
 		return r.deleteHeadlessServiceIfExists(ctx, lws, lws.Name)

diff --git a/test/e2e/e2e_test.go b/test/e2e/e2e_test.go
@@ -209,19 +209,28 @@ var _ = ginkgo.Describe("leaderWorkerSet e2e tests", func() {
 		testing.MustCreateLws(ctx, k8sClient, lws)
 		testing.ExpectValidPods(ctx, k8sClient, lws, &corev1.PodList{})
 		testing.UpdateSubdomainPolicy(ctx, k8sClient, lws, leaderworkerset.SubdomainUniquePerReplica)
-		testing.ExpectValidPods(ctx, k8sClient, lws, &corev1.PodList{})
-		testing.UpdateReplicaCount(ctx, k8sClient, lws, 2)
-
 		lwsPods := &corev1.PodList{}
 		testing.ExpectValidPods(ctx, k8sClient, lws, lwsPods)
 
 		for _, pod := range lwsPods.Items {
+			gomega.Expect(testing.HasTPUEnvVarsPopulated(pod)).To(gomega.BeTrue())
 			if pod.Labels[leaderworkerset.GroupIndexLabelKey] == "0" {
 				gomega.Expect(testing.CheckTPUContainerHasCorrectEnvVars(pod, "test-sample-0.test-sample-0,test-sample-0-1.test-sample-0")).Should(gomega.Succeed())
 			} else {
 				gomega.Expect(testing.CheckTPUContainerHasCorrectEnvVars(pod, "test-sample-1.test-sample-1,test-sample-1-1.test-sample-1")).Should(gomega.Succeed())
 			}
+		}
+		testing.UpdateSubdomainPolicy(ctx, k8sClient, lws, leaderworkerset.SubdomainShared)
+		lwsPodsAfterUpgrade := &corev1.PodList{}
+		testing.ExpectValidPods(ctx, k8sClient, lws, lwsPodsAfterUpgrade)
+
+		for _, pod := range lwsPodsAfterUpgrade.Items {
 			gomega.Expect(testing.HasTPUEnvVarsPopulated(pod)).To(gomega.BeTrue())
+			if pod.Labels[leaderworkerset.GroupIndexLabelKey] == "0" {
+				gomega.Expect(testing.CheckTPUContainerHasCorrectEnvVars(pod, "test-sample-0.test-sample,test-sample-0-1.test-sample")).Should(gomega.Succeed())
+			} else {
+				gomega.Expect(testing.CheckTPUContainerHasCorrectEnvVars(pod, "test-sample-1.test-sample,test-sample-1-1.test-sample")).Should(gomega.Succeed())
+			}
 		}
 	})