[RayService] Add unit tests for isZeroDowntimeUpgradeEnabled (#2871)

Author: kevin85421

ray-operator/controllers/ray/rayservice_controller.go

@@ -463,14 +463,13 @@ func (r *RayServiceReconciler) getRayServiceInstance(ctx context.Context, reques
 	return rayServiceInstance, nil
 }
 
-func isZeroDowntimeUpgradeEnabled(ctx context.Context, rayService *rayv1.RayService) bool {
+func isZeroDowntimeUpgradeEnabled(ctx context.Context, upgradeStrategy *rayv1.RayServiceUpgradeStrategy) bool {
 	// For LLM serving, some users might not have sufficient GPU resources to run two RayClusters simultaneously.
 	// Therefore, KubeRay offers ENABLE_ZERO_DOWNTIME as a feature flag for zero-downtime upgrades.
 	// There are two ways to enable zero downtime upgrade. Through ENABLE_ZERO_DOWNTIME env var or setting Spec.UpgradeStrategy.Type.
 	// If no fields are set, zero downtime upgrade by default is enabled.
 	// Spec.UpgradeStrategy.Type takes precedence over ENABLE_ZERO_DOWNTIME.
 	logger := ctrl.LoggerFrom(ctx)
-	upgradeStrategy := rayService.Spec.UpgradeStrategy
 	if upgradeStrategy != nil {
 		upgradeType := upgradeStrategy.Type
 		if upgradeType != nil {
@@ -691,7 +690,7 @@ func shouldPrepareNewCluster(ctx context.Context, rayServiceInstance *rayv1.RayS
 		// KubeRay should update the RayCluster instead of creating a new one.
 		return false
 	}
-	return isZeroDowntimeUpgradeEnabled(ctx, rayServiceInstance)
+	return isZeroDowntimeUpgradeEnabled(ctx, rayServiceInstance.Spec.UpgradeStrategy)
 }
 
 // updateRayClusterInstance updates the RayCluster instance.

ray-operator/controllers/ray/rayservice_controller_test.go

@@ -18,7 +18,6 @@ package ray
 import (
 	"context"
 	"fmt"
-	"os"
 	"time"
 
 	"github.com/ray-project/kuberay/ray-operator/controllers/ray/common"
@@ -344,47 +343,6 @@ var _ = Context("RayService env tests", func() {
 			Expect(meta.IsStatusConditionFalse(rayService.Status.Conditions, string(rayv1.UpgradeInProgress))).Should(BeTrue())
 		})
 
-		It("Disable zero-downtime upgrade", func() {
-			// Disable zero-downtime upgrade.
-			os.Setenv("ENABLE_ZERO_DOWNTIME", "false")
-
-			// Try to trigger a zero-downtime upgrade.
-			oldRayVersion := rayService.Spec.RayClusterSpec.RayVersion
-			newRayVersion := "2.198.0"
-			Expect(oldRayVersion).ShouldNot(Equal(newRayVersion))
-			err := retry.RetryOnConflict(retry.DefaultRetry, func() error {
-				Eventually(
-					getResourceFunc(ctx, client.ObjectKey{Name: rayService.Name, Namespace: "default"}, rayService),
-					time.Second*3, time.Millisecond*500).Should(BeNil(), "My myRayService  = %v", rayService.Name)
-				rayService.Spec.RayClusterSpec.RayVersion = newRayVersion
-				return k8sClient.Update(ctx, rayService)
-			})
-			Expect(err).NotTo(HaveOccurred(), "failed to update test RayService resource")
-
-			// Because the zero-downtime upgrade is disabled, the RayService controller will not prepare a new RayCluster.
-			Consistently(
-				getPreparingRayClusterNameFunc(ctx, rayService),
-				time.Second*5, time.Millisecond*500).Should(BeEmpty(), "Pending RayCluster name  = %v", rayService.Status.PendingServiceStatus.RayClusterName)
-
-			// Set the RayVersion back to the old value to avoid triggering the zero-downtime upgrade.
-			err = retry.RetryOnConflict(retry.DefaultRetry, func() error {
-				Eventually(
-					getResourceFunc(ctx, client.ObjectKey{Name: rayService.Name, Namespace: "default"}, rayService),
-					time.Second*3, time.Millisecond*500).Should(BeNil(), "My myRayService  = %v", rayService.Name)
-				rayService.Spec.RayClusterSpec.RayVersion = oldRayVersion
-				return k8sClient.Update(ctx, rayService)
-			})
-			Expect(err).NotTo(HaveOccurred(), "failed to update test RayService resource")
-
-			// Enable zero-downtime upgrade again.
-			os.Unsetenv("ENABLE_ZERO_DOWNTIME")
-
-			// Zero-downtime upgrade should not be triggered.
-			Consistently(
-				getPreparingRayClusterNameFunc(ctx, rayService),
-				time.Second*5, time.Millisecond*500).Should(BeEmpty(), "Pending RayCluster name  = %v", rayService.Status.PendingServiceStatus.RayClusterName)
-		})
-
 		It("Autoscaler updates the active RayCluster and should not switch to a new RayCluster", func() {
 			// Simulate autoscaler by updating the active RayCluster directly. Note that the autoscaler
 			// will not update the RayService directly.

ray-operator/controllers/ray/rayservice_controller_unit_test.go

@@ -3,6 +3,7 @@ package ray
 import (
 	"context"
 	"fmt"
+	"os"
 	"reflect"
 	"strconv"
 	"testing"
@@ -1155,3 +1156,81 @@ func TestShouldPrepareNewCluster_ZeroDowntimeUpgrade(t *testing.T) {
 	shouldPrepareNewCluster := shouldPrepareNewCluster(ctx, &rayService, activeCluster, nil)
 	assert.True(t, shouldPrepareNewCluster)
 }
+
+func TestIsZeroDowntimeUpgradeEnabled(t *testing.T) {
+	tests := []struct {
+		name                     string
+		upgradeStrategy          *rayv1.RayServiceUpgradeStrategy
+		enableZeroDowntimeEnvVar string // "true" or "false" or "" (not set)
+		expected                 bool
+	}{
+		{
+			// The most common case.
+			name:                     "both upgrade strategy and env var are not set",
+			upgradeStrategy:          nil,
+			enableZeroDowntimeEnvVar: "",
+			expected:                 true,
+		},
+		{
+			name:                     "upgrade strategy is not set, but env var is set to true",
+			upgradeStrategy:          nil,
+			enableZeroDowntimeEnvVar: "true",
+			expected:                 true,
+		},
+		{
+			name:                     "upgrade strategy is not set, but env var is set to false",
+			upgradeStrategy:          nil,
+			enableZeroDowntimeEnvVar: "false",
+			expected:                 false,
+		},
+		{
+			name:                     "upgrade strategy is set to NewCluster",
+			upgradeStrategy:          &rayv1.RayServiceUpgradeStrategy{Type: ptr.To(rayv1.NewCluster)},
+			enableZeroDowntimeEnvVar: "",
+			expected:                 true,
+		},
+		{
+			name:                     "upgrade strategy is set to NewCluster, and env var is not set",
+			upgradeStrategy:          &rayv1.RayServiceUpgradeStrategy{Type: ptr.To(rayv1.NewCluster)},
+			enableZeroDowntimeEnvVar: "true",
+			expected:                 true,
+		},
+		{
+			name:                     "upgrade strategy is set to NewCluster, and env var is set to false",
+			upgradeStrategy:          &rayv1.RayServiceUpgradeStrategy{Type: ptr.To(rayv1.NewCluster)},
+			enableZeroDowntimeEnvVar: "false",
+			expected:                 true,
+		},
+		{
+			name:                     "upgrade strategy is set to None, and env var is not set",
+			upgradeStrategy:          &rayv1.RayServiceUpgradeStrategy{Type: ptr.To(rayv1.None)},
+			enableZeroDowntimeEnvVar: "",
+			expected:                 false,
+		},
+		{
+			name:                     "upgrade strategy is set to None, and env var is set to true",
+			upgradeStrategy:          &rayv1.RayServiceUpgradeStrategy{Type: ptr.To(rayv1.None)},
+			enableZeroDowntimeEnvVar: "true",
+			expected:                 false,
+		},
+		{
+			name:                     "upgrade strategy is set to None, and env var is set to false",
+			upgradeStrategy:          &rayv1.RayServiceUpgradeStrategy{Type: ptr.To(rayv1.None)},
+			enableZeroDowntimeEnvVar: "false",
+			expected:                 false,
+		},
+	}
+
+	ctx := context.TODO()
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			defer func() {
+				os.Unsetenv(ENABLE_ZERO_DOWNTIME)
+			}()
+
+			os.Setenv(ENABLE_ZERO_DOWNTIME, tt.enableZeroDowntimeEnvVar)
+			isEnabled := isZeroDowntimeUpgradeEnabled(ctx, tt.upgradeStrategy)
+			assert.Equal(t, tt.expected, isEnabled)
+		})
+	}
+}