[Refactor] Use RAYCLUSTER_DEFAULT_REQUEUE_SECONDS_ENV as timeout of status check in tests

Author: rueian

ray-operator/controllers/ray/raycluster_controller.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"os"
 	"reflect"
-	"strconv"
 	"strings"
 	"time"
 
@@ -382,10 +381,9 @@ func (r *RayClusterReconciler) rayClusterReconcile(ctx context.Context, request
 	// Unconditionally requeue after the number of seconds specified in the
 	// environment variable RAYCLUSTER_DEFAULT_REQUEUE_SECONDS_ENV. If the
 	// environment variable is not set, requeue after the default value.
-	requeueAfterSeconds, err := strconv.Atoi(os.Getenv(utils.RAYCLUSTER_DEFAULT_REQUEUE_SECONDS_ENV))
-	if err != nil {
+	requeueAfterSeconds, fellback := utils.GetEnvInt(utils.RAYCLUSTER_DEFAULT_REQUEUE_SECONDS_ENV, utils.RAYCLUSTER_DEFAULT_REQUEUE_SECONDS)
+	if fellback {
 		r.Log.Info(fmt.Sprintf("Environment variable %s is not set, using default value of %d seconds", utils.RAYCLUSTER_DEFAULT_REQUEUE_SECONDS_ENV, utils.RAYCLUSTER_DEFAULT_REQUEUE_SECONDS), "cluster name", request.Name)
-		requeueAfterSeconds = utils.RAYCLUSTER_DEFAULT_REQUEUE_SECONDS
 	}
 	r.Log.Info("Unconditional requeue after", "cluster name", request.Name, "seconds", requeueAfterSeconds)
 	return ctrl.Result{RequeueAfter: time.Duration(requeueAfterSeconds) * time.Second}, nil

ray-operator/controllers/ray/raycluster_controller_test.go

@@ -205,7 +205,7 @@ var _ = Context("Inside the default namespace", func() {
 			// We need to figure out the behavior. See https://github.com/ray-project/kuberay/issues/1736 for more details.
 			Eventually(
 				getClusterState(ctx, "default", myRayCluster.Name),
-				time.Second*(utils.RAYCLUSTER_DEFAULT_REQUEUE_SECONDS+5), time.Millisecond*500).Should(Equal(rayv1.Ready))
+				time.Second*(time.Duration(utils.MustGetEnvInt(utils.RAYCLUSTER_DEFAULT_REQUEUE_SECONDS_ENV)+5)), time.Millisecond*500).Should(Equal(rayv1.Ready))
 		})
 
 		It("should re-create a deleted worker", func() {
@@ -311,7 +311,7 @@ var _ = Context("Inside the default namespace", func() {
 		It("cluster's .status.state should be updated to 'suspended' shortly after all Pods are terminated", func() {
 			Eventually(
 				getClusterState(ctx, "default", myRayCluster.Name),
-				time.Second*(utils.RAYCLUSTER_DEFAULT_REQUEUE_SECONDS+5), time.Millisecond*500).Should(Equal(rayv1.Suspended))
+				time.Second*(time.Duration(utils.MustGetEnvInt(utils.RAYCLUSTER_DEFAULT_REQUEUE_SECONDS_ENV)+5)), time.Millisecond*500).Should(Equal(rayv1.Suspended))
 		})
 
 		It("set suspend to false and then revert it to true before all Pods are running", func() {
@@ -363,7 +363,7 @@ var _ = Context("Inside the default namespace", func() {
 			// RayCluster should be in Suspended state.
 			Eventually(
 				getClusterState(ctx, "default", myRayCluster.Name),
-				time.Second*(utils.RAYCLUSTER_DEFAULT_REQUEUE_SECONDS+5), time.Millisecond*500).Should(Equal(rayv1.Suspended))
+				time.Second*(time.Duration(utils.MustGetEnvInt(utils.RAYCLUSTER_DEFAULT_REQUEUE_SECONDS_ENV)+5)), time.Millisecond*500).Should(Equal(rayv1.Suspended))
 		})
 
 		It("should run all head and worker pods if un-suspended", func() {
@@ -402,7 +402,7 @@ var _ = Context("Inside the default namespace", func() {
 		It("cluster's .status.state should be updated back to 'ready' after being un-suspended", func() {
 			Eventually(
 				getClusterState(ctx, "default", myRayCluster.Name),
-				time.Second*(utils.RAYCLUSTER_DEFAULT_REQUEUE_SECONDS+5), time.Millisecond*500).Should(Equal(rayv1.Ready))
+				time.Second*(time.Duration(utils.MustGetEnvInt(utils.RAYCLUSTER_DEFAULT_REQUEUE_SECONDS_ENV)+5)), time.Millisecond*500).Should(Equal(rayv1.Ready))
 		})
 	})
 })

ray-operator/controllers/ray/rayjob_controller_test.go

@@ -270,7 +270,7 @@ var _ = Context("Inside the default namespace", func() {
 			// The RayCluster.Status.State should be Ready.
 			Eventually(
 				getClusterState(ctx, "default", myRayCluster.Name),
-				time.Second*(utils.RAYCLUSTER_DEFAULT_REQUEUE_SECONDS+5), time.Millisecond*500).Should(Equal(rayv1.Ready))
+				time.Second*(time.Duration(utils.MustGetEnvInt(utils.RAYCLUSTER_DEFAULT_REQUEUE_SECONDS_ENV)+5)), time.Millisecond*500).Should(Equal(rayv1.Ready))
 		})
 
 		It("Dashboard URL should be set", func() {

ray-operator/controllers/ray/utils/constant.go

@@ -1,5 +1,11 @@
 package utils
 
+import (
+	"fmt"
+	"os"
+	"strconv"
+)
+
 const (
 
 	// Default application name
@@ -143,3 +149,19 @@ const (
 	HeadService    ServiceType = "headService"
 	ServingService ServiceType = "serveService"
 )
+
+func GetEnvInt(name string, fallback int) (value int, fellback bool) {
+	v, err := strconv.Atoi(os.Getenv(name))
+	if err != nil {
+		v = fallback
+	}
+	return v, err != nil
+}
+
+func MustGetEnvInt(name string) (value int) {
+	v, err := strconv.Atoi(os.Getenv(name))
+	if err != nil {
+		panic(fmt.Sprintf("env %s must be set: %v", name, err))
+	}
+	return v
+}