feat: Add generics method to check status of sts svc (#2)

Signed-off-by: chandankumar4 <[email protected]>
numaproj-contrib · Oct 31, 2023 · ebb98ab · ebb98ab
1 parent e2700cf
commit ebb98ab
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 7 deletions.
diff --git a/testing/fixtures/util.go b/testing/fixtures/util.go
@@ -124,8 +124,7 @@ func WaitForISBSvcReady(ctx context.Context, isbSvcClient flowpkg.InterStepBuffe
 	}
 }
 
-func WaitForISBSvcStatefulSetReady(ctx context.Context, kubeClient kubernetes.Interface, namespace, isbSvcName string, timeout time.Duration) error {
-	labelSelector := fmt.Sprintf("%s=isbsvc-controller,%s=%s", dfv1.KeyManagedBy, dfv1.KeyISBSvcName, isbSvcName)
+func WaitForStatefulSetReady(ctx context.Context, kubeClient kubernetes.Interface, timeout time.Duration, namespace, labelSelector string) error {
 	opts := metav1.ListOptions{LabelSelector: labelSelector}
 	watch, err := kubeClient.AppsV1().StatefulSets(namespace).Watch(ctx, opts)
 	if err != nil {
@@ -138,20 +137,27 @@ func WaitForISBSvcStatefulSetReady(ctx context.Context, kubeClient kubernetes.In
 		timeoutCh <- true
 	}()
 
+	var (
+		replicas int32
+		svcName  string
+	)
+
 statefulSetWatch:
 	for {
 		select {
 		case event := <-watch.ResultChan():
 			ss, ok := event.Object.(*appsv1.StatefulSet)
 			if ok {
 				if ss.Status.Replicas == ss.Status.ReadyReplicas {
+					replicas = ss.Status.Replicas
+					svcName = ss.Name
 					break statefulSetWatch
 				}
 			} else {
 				return fmt.Errorf("not statefulset")
 			}
 		case <-timeoutCh:
-			return fmt.Errorf("timeout after %v waiting for ISB svc StatefulSet ready", timeout)
+			return fmt.Errorf("timeout after %v waiting for %s svc StatefulSet ready", timeout, svcName)
 		}
 	}
 
@@ -169,8 +175,7 @@ statefulSetWatch:
 
 	podNames := make(map[string]bool)
 	for {
-		if len(podNames) == 3 {
-			// defaults to 3 Pods
+		if len(podNames) == int(replicas) {
 			return nil
 		}
 		select {
@@ -194,7 +199,7 @@ statefulSetWatch:
 				return fmt.Errorf("not pod")
 			}
 		case <-podTimeoutCh:
-			return fmt.Errorf("timeout after %v waiting for ISB svc Pod ready", timeout)
+			return fmt.Errorf("timeout after %v waiting for %s Pod ready", timeout, svcName)
 		}
 	}
 }

diff --git a/testing/fixtures/when.go b/testing/fixtures/when.go
@@ -159,7 +159,19 @@ func (w *When) WaitForISBSvcReady() *When {
 	if err := WaitForISBSvcReady(ctx, w.isbSvcClient, w.isbSvc.Name, defaultTimeout); err != nil {
 		w.t.Fatal(err)
 	}
-	if err := WaitForISBSvcStatefulSetReady(ctx, w.kubeClient, Namespace, w.isbSvc.Name, 5*time.Minute); err != nil {
+
+	isbSvclabels := fmt.Sprintf("%s=isbsvc-controller,%s=%s", dfv1.KeyManagedBy, dfv1.KeyISBSvcName, w.isbSvc.Name)
+	if err := w.WaitForStatefulSetReady(isbSvclabels); err != nil {
+		w.t.Fatal(err)
+	}
+
+	return w
+}
+
+func (w *When) WaitForStatefulSetReady(labelSelector string) *When {
+	w.t.Helper()
+	ctx := context.Background()
+	if err := WaitForStatefulSetReady(ctx, w.kubeClient, 5*time.Minute, Namespace, labelSelector); err != nil {
 		w.t.Fatal(err)
 	}
 	return w