[Apiserver] Use Eventually from Gomega instead of wait from apimachinery (#3433)

Author: owenowenisme

apiserver/go.mod

@@ -27,6 +27,7 @@ require (
 	github.com/grpc-ecosystem/go-grpc-middleware v1.4.0
 	github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0
 	github.com/grpc-ecosystem/grpc-gateway/v2 v2.20.0
+	github.com/onsi/gomega v1.36.2
 	github.com/rs/zerolog v1.34.0
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20240826202546-f6391c0de4c7
 	k8s.io/utils v0.0.0-20241210054802-24370beab758

apiserver/test/e2e/cluster_server_autoscaler_e2e_test.go

@@ -1,21 +1,21 @@
 package e2e
 
 import (
-	"context"
 	"testing"
-	"time"
-
-	"github.com/stretchr/testify/require"
-	"k8s.io/apimachinery/pkg/util/wait"
 
 	api "github.com/ray-project/kuberay/proto/go_client"
 
 	rayv1api "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
+
+	"github.com/onsi/gomega"
+	"github.com/stretchr/testify/require"
 )
 
 // TestCreateClusterAutoscalerEndpoint sequentially iterates over the create cluster endpoint
 // with valid and invalid requests
 func TestCreateClusterAutoscaler(t *testing.T) {
+	g := gomega.NewWithT(t)
+
 	tCtx, err := NewEnd2EndTestingContext(t)
 	require.NoError(t, err, "No error expected when creating testing context")
 
@@ -120,18 +120,18 @@ func TestCreateClusterAutoscaler(t *testing.T) {
 	require.NoError(t, err, "No error expected")
 	require.Nil(t, actualRPCStatus, "No RPC status expected")
 	require.NotNil(t, actualJob, "A job is expected")
-	waitForRayJob(t, tCtx, createActorRequest.Job.Name, []rayv1api.JobStatus{rayv1api.JobStatusSucceeded})
+	waitForRayJobInExpectedStatuses(t, tCtx, createActorRequest.Job.Name, []rayv1api.JobStatus{rayv1api.JobStatusSucceeded})
 
 	// worker pod should be created as part of job execution
-	err = wait.PollUntilContextTimeout(tCtx.ctx, 500*time.Millisecond, 3*time.Minute, false, func(_ context.Context) (done bool, err error) {
+	g.Eventually(func() int32 {
 		rayCluster, err := tCtx.GetRayClusterByName(actualCluster.Name)
 		if err != nil {
-			return true, err
+			t.Logf("Error getting ray cluster: %v", err)
+			return -1 // Return -1 to indicate error condition
 		}
 		t.Logf("Found ray cluster with %d available worker replicas", rayCluster.Status.AvailableWorkerReplicas)
-		return rayCluster.Status.AvailableWorkerReplicas == 1, nil
-	})
-	require.NoError(t, err, "No error expected")
+		return rayCluster.Status.AvailableWorkerReplicas
+	}, TestTimeoutMedium, TestPollingInterval).Should(gomega.Equal(int32(1)))
 	// Delete actor
 	deleteActorRequest := api.CreateRayJobRequest{
 		Namespace: tCtx.GetNamespaceName(),
@@ -149,15 +149,15 @@ func TestCreateClusterAutoscaler(t *testing.T) {
 	require.NoError(t, err, "No error expected")
 	require.Nil(t, actualRPCStatus, "No RPC status expected")
 	require.NotNil(t, actualJob, "A job is expected")
-	waitForRayJob(t, tCtx, createActorRequest.Job.Name, []rayv1api.JobStatus{rayv1api.JobStatusSucceeded})
-
-	err = wait.PollUntilContextTimeout(tCtx.ctx, 500*time.Millisecond, 3*time.Minute, false, func(_ context.Context) (done bool, err error) {
+	waitForRayJobInExpectedStatuses(t, tCtx, createActorRequest.Job.Name, []rayv1api.JobStatus{rayv1api.JobStatusSucceeded})
+	g.Eventually(func() int32 {
 		rayCluster, err := tCtx.GetRayClusterByName(actualCluster.Name)
 		if err != nil {
-			return true, err
+			t.Logf("Error getting ray cluster: %v", err)
+			return -1 // Return -1 to indicate error condition
 		}
 		t.Logf("Found ray cluster with %d available worker replicas", rayCluster.Status.AvailableWorkerReplicas)
-		return rayCluster.Status.AvailableWorkerReplicas == 0, nil
-	})
+		return rayCluster.Status.AvailableWorkerReplicas
+	}, TestTimeoutMedium, TestPollingInterval).Should(gomega.Equal(int32(0)))
 	require.NoError(t, err, "No error expected")
 }

apiserver/test/e2e/cluster_server_e2e_test.go

@@ -5,12 +5,12 @@ import (
 	"net/http"
 	"testing"
 
-	"github.com/stretchr/testify/require"
-
 	kuberayHTTP "github.com/ray-project/kuberay/apiserver/pkg/http"
 	api "github.com/ray-project/kuberay/proto/go_client"
 
 	rayv1api "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
+
+	"github.com/stretchr/testify/require"
 )
 
 // TestCreateClusterEndpoint sequentially iterates over the create cluster endpoint
@@ -503,7 +503,7 @@ func TestDeleteCluster(t *testing.T) {
 			if tc.ExpectedError == nil {
 				require.NoError(t, err, "No error expected")
 				require.Nil(t, actualRPCStatus, "No RPC status expected")
-				waitForDeletedCluster(t, tCtx, tc.Input.Name)
+				waitForClusterToDisappear(t, tCtx, tc.Input.Name)
 			} else {
 				require.EqualError(t, err, tc.ExpectedError.Error(), "Matching error expected")
 				require.NotNil(t, actualRPCStatus, "A not nill RPC status is required")

apiserver/test/e2e/config_server_e2e_test.go

@@ -5,10 +5,10 @@ import (
 	"reflect"
 	"testing"
 
-	"github.com/stretchr/testify/require"
-
 	kuberayHTTP "github.com/ray-project/kuberay/apiserver/pkg/http"
 	api "github.com/ray-project/kuberay/proto/go_client"
+
+	"github.com/stretchr/testify/require"
 )
 
 // TestCreateTemplate sequentially iterates over the create compute endpoint

apiserver/test/e2e/job_server_e2e_test.go

@@ -1,21 +1,16 @@
 package e2e
 
 import (
-	"context"
 	"fmt"
 	"net/http"
-	"slices"
 	"testing"
-	"time"
-
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-	"k8s.io/apimachinery/pkg/util/wait"
 
 	kuberayHTTP "github.com/ray-project/kuberay/apiserver/pkg/http"
 	api "github.com/ray-project/kuberay/proto/go_client"
 
 	rayv1api "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestCreateJobWithDisposableClusters(t *testing.T) {
@@ -94,7 +89,6 @@ func TestCreateJobWithDisposableClusters(t *testing.T) {
 					ShutdownAfterJobFinishes: true,
 					ClusterSpec:              clusterSpec,
 					TtlSecondsAfterFinished:  60,
-					ActiveDeadlineSeconds:    120,
 				},
 				Namespace: tCtx.GetNamespaceName(),
 			},
@@ -113,7 +107,6 @@ func TestCreateJobWithDisposableClusters(t *testing.T) {
 					ShutdownAfterJobFinishes: true,
 					ClusterSpec:              clusterSpec,
 					TtlSecondsAfterFinished:  60,
-					ActiveDeadlineSeconds:    120,
 				},
 				Namespace: tCtx.GetNamespaceName(),
 			},
@@ -223,7 +216,7 @@ func TestCreateJobWithDisposableClusters(t *testing.T) {
 				require.NoError(t, err, "No error expected")
 				require.Nil(t, actualRPCStatus, "No RPC status expected")
 				require.NotNil(t, actualJob, "A job is expected")
-				waitForRayJob(t, tCtx, tc.Input.Job.Name, []rayv1api.JobStatus{tc.ExpectedJobStatus})
+				waitForRayJobInExpectedStatuses(t, tCtx, tc.Input.Job.Name, []rayv1api.JobStatus{tc.ExpectedJobStatus})
 				tCtx.DeleteRayJobByName(t, actualJob.Name)
 			} else {
 				require.EqualError(t, err, tc.ExpectedError.Error(), "Matching error expected")
@@ -281,7 +274,7 @@ func TestDeleteJob(t *testing.T) {
 			if tc.ExpectedError == nil {
 				require.NoError(t, err, "No error expected")
 				require.Nil(t, actualRPCStatus, "No RPC status expected")
-				waitForDeletedRayJob(t, tCtx, testJobRequest.Job.Name)
+				waitForRayJobToDisappear(t, tCtx, testJobRequest.Job.Name)
 			} else {
 				require.EqualError(t, err, tc.ExpectedError.Error(), "Matching error expected")
 				require.NotNil(t, actualRPCStatus, "A not nill RPC status is required")
@@ -650,7 +643,7 @@ func TestCreateJobWithClusterSelector(t *testing.T) {
 				require.NoError(t, err, "No error expected")
 				require.Nil(t, actualRPCStatus, "No RPC status expected")
 				require.NotNil(t, actualJob, "A job is expected")
-				waitForRayJob(t, tCtx, tc.Input.Job.Name, []rayv1api.JobStatus{tc.ExpectedJobStatus})
+				waitForRayJobInExpectedStatuses(t, tCtx, tc.Input.Job.Name, []rayv1api.JobStatus{tc.ExpectedJobStatus})
 				tCtx.DeleteRayJobByName(t, actualJob.Name)
 			} else {
 				require.EqualError(t, err, tc.ExpectedError.Error(), "Matching error expected")
@@ -728,36 +721,6 @@ func createTestJob(t *testing.T, tCtx *End2EndTestingContext, expectedJobStatues
 	require.NoError(t, err, "No error expected")
 	require.Nil(t, actualRPCStatus, "No RPC status expected")
 	require.NotNil(t, actualJob, "A job is expected")
-	waitForRayJob(t, tCtx, testJobRequest.Job.Name, expectedJobStatues)
+	waitForRayJobInExpectedStatuses(t, tCtx, testJobRequest.Job.Name, expectedJobStatues)
 	return testJobRequest
 }
-
-func waitForRayJob(t *testing.T, tCtx *End2EndTestingContext, rayJobName string, expectedJobStatuses []rayv1api.JobStatus) {
-	// `expectedJobStatuses` is a slice of job statuses that we expect the job to be in
-	// wait for the job to be in any of the `expectedJobStatuses` state for 3 minutes
-	// if is not in that state, return an error
-	err := wait.PollUntilContextTimeout(tCtx.ctx, 500*time.Millisecond, 3*time.Minute, false, func(_ context.Context) (done bool, err error) {
-		rayJob, err := tCtx.GetRayJobByName(rayJobName)
-		if err != nil {
-			return true, err
-		}
-		t.Logf("Found ray job with state '%s' for ray job '%s'", rayJob.Status.JobStatus, rayJobName)
-		return slices.Contains(expectedJobStatuses, rayJob.Status.JobStatus), nil
-	})
-	require.NoErrorf(t, err, "No error expected when getting status for ray job: '%s', err %v", rayJobName, err)
-}
-
-func waitForDeletedRayJob(t *testing.T, tCtx *End2EndTestingContext, jobName string) {
-	// wait for the job to be deleted
-	// if is not in that state, return an error
-	err := wait.PollUntilContextTimeout(tCtx.ctx, 500*time.Millisecond, 3*time.Minute, false, func(_ context.Context) (done bool, err error) {
-		rayJob, err := tCtx.GetRayJobByName(jobName)
-		if err != nil &&
-			assert.EqualError(t, err, "rayjobs.ray.io \""+jobName+"\" not found") {
-			return true, nil
-		}
-		t.Logf("Found status of '%s' for ray cluster '%s'", rayJob.Status.JobStatus, jobName)
-		return false, err
-	})
-	require.NoErrorf(t, err, "No error expected when deleting ray job: '%s', err %v", jobName, err)
-}

apiserver/test/e2e/job_submission_e2e_test.go

@@ -3,9 +3,9 @@ package e2e
 import (
 	"testing"
 
-	"github.com/stretchr/testify/require"
-
 	api "github.com/ray-project/kuberay/proto/go_client"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestCreateJobSubmission(t *testing.T) {

apiserver/test/e2e/service_server_e2e_test.go

@@ -1,18 +1,13 @@
 package e2e
 
 import (
-	"context"
 	"net/http"
 	"testing"
-	"time"
-
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-	"google.golang.org/protobuf/proto"
-	"k8s.io/apimachinery/pkg/util/wait"
 
 	kuberayHTTP "github.com/ray-project/kuberay/apiserver/pkg/http"
 	api "github.com/ray-project/kuberay/proto/go_client"
+
+	"github.com/stretchr/testify/require"
 )
 
 // TestServiceServerV2 sequentially iterates over the endpoints of the service endpoints using
@@ -199,7 +194,7 @@ func TestDeleteService(t *testing.T) {
 			if tc.ExpectedError == nil {
 				require.NoError(t, err, "No error expected")
 				require.Nil(t, actualRPCStatus, "No RPC status expected")
-				waitForDeletedService(t, tCtx, testServiceRequest.Service.Name)
+				waitForServiceToDisappear(t, tCtx, testServiceRequest.Service.Name)
 			} else {
 				require.EqualError(t, err, tc.ExpectedError.Error(), "Matching error expected")
 				require.NotNil(t, actualRPCStatus, "A not nill RPC status is required")
@@ -208,66 +203,6 @@ func TestDeleteService(t *testing.T) {
 	}
 }
 
-func TestUpdateRayService(t *testing.T) {
-	tCtx, err := NewEnd2EndTestingContext(t)
-	require.NoError(t, err)
-
-	tCtx.CreateComputeTemplate(t)
-	t.Cleanup(func() {
-		tCtx.DeleteComputeTemplate(t)
-	})
-
-	testServiceRequest := createTestServiceV2(t, tCtx)
-	serviceName := testServiceRequest.Service.Name
-	namespace := testServiceRequest.Namespace
-
-	t.Cleanup(func() {
-		tCtx.DeleteRayService(t, serviceName)
-	})
-
-	// Verify the original value before update
-	require.Equal(t, int32(1), testServiceRequest.Service.ClusterSpec.WorkerGroupSpec[0].Replicas)
-
-	// Clone the original ClusterSpec (proto.Clone avoids copying sync.Mutex)
-	oldSpec := proto.Clone(testServiceRequest.Service.ClusterSpec).(*api.ClusterSpec)
-
-	// Only modify replicas
-	oldSpec.WorkerGroupSpec[0].Replicas = 2 // Only field updated in this test
-
-	// Construct updated service from old spec
-	updatedService := &api.RayService{
-		Name:                               serviceName,
-		Namespace:                          namespace,
-		User:                               testServiceRequest.Service.User,
-		Version:                            testServiceRequest.Service.Version,
-		ServeConfig_V2:                     testServiceRequest.Service.ServeConfig_V2,
-		ServiceUnhealthySecondThreshold:    testServiceRequest.Service.ServiceUnhealthySecondThreshold,
-		DeploymentUnhealthySecondThreshold: testServiceRequest.Service.DeploymentUnhealthySecondThreshold,
-		ClusterSpec:                        oldSpec,
-	}
-
-	updateReq := &api.UpdateRayServiceRequest{
-		Service:   updatedService,
-		Namespace: namespace,
-		Name:      serviceName,
-	}
-
-	// Perform the update
-	respService, actualRPCStatus, err := tCtx.GetRayAPIServerClient().UpdateRayService(updateReq)
-	require.NoError(t, err)
-	require.Nil(t, actualRPCStatus)
-	require.NotNil(t, respService)
-
-	// Confirm update via RayService CRD
-	crdRayService, err := tCtx.GetRayServiceByName(serviceName)
-	require.NoError(t, err)
-	require.NotNil(t, crdRayService)
-
-	require.GreaterOrEqual(t, len(crdRayService.Spec.RayClusterSpec.WorkerGroupSpecs), 1)
-	require.NotNil(t, crdRayService.Spec.RayClusterSpec.WorkerGroupSpecs[0].Replicas)
-	require.Equal(t, int32(2), *crdRayService.Spec.RayClusterSpec.WorkerGroupSpecs[0].Replicas)
-}
-
 func TestGetAllServices(t *testing.T) {
 	tCtx, err := NewEnd2EndTestingContext(t)
 	require.NoError(t, err, "No error expected when creating testing context")
@@ -537,18 +472,3 @@ func checkRayServiceCreatedSuccessfully(t *testing.T, tCtx *End2EndTestingContex
 	require.NoError(t, err)
 	require.NotNil(t, rayService)
 }
-
-func waitForDeletedService(t *testing.T, tCtx *End2EndTestingContext, serviceName string) {
-	// wait for the service to be deleted
-	// if is not in that state, return an error
-	err := wait.PollUntilContextTimeout(tCtx.ctx, 500*time.Millisecond, 3*time.Minute, false, func(_ context.Context) (done bool, err error) {
-		rayService, err := tCtx.GetRayServiceByName(serviceName)
-		if err != nil &&
-			assert.EqualError(t, err, "rayservices.ray.io \""+serviceName+"\" not found") {
-			return true, nil
-		}
-		t.Logf("Found status of '%s' for ray service '%s'", rayService.Status.ServiceStatus, serviceName)
-		return false, err
-	})
-	require.NoErrorf(t, err, "No error expected when deleting ray service: '%s', err %v", serviceName, err)
-}