test: enable upgrade to image built from source (#3736)

Author: pawelpaszki

ray-operator/test/e2eupgrade/rayservice_operator_upgrade_test.go

@@ -2,6 +2,7 @@ package e2eupgrade
 
 import (
 	"fmt"
+	"os"
 	"os/exec"
 	"testing"
 
@@ -33,6 +34,8 @@ func TestZeroDowntimeUpgradeAfterOperatorUpgrade(t *testing.T) {
 	// Get the upgrade version from environment
 	upgradeVersion := GetKubeRayUpgradeVersion()
 
+	test.T().Logf("Detected upgrade version: %s", upgradeVersion)
+
 	// Create RayService custom resource
 	rayServiceAC := rayv1ac.RayService(rayServiceName, namespace.Name).WithSpec(e2e.RayServiceSampleYamlApplyConfiguration())
 	rayService, err := test.Client().Ray().RayV1().RayServices(namespace.Name).Apply(test.Ctx(), rayServiceAC, TestApplyOptions)
@@ -76,15 +79,38 @@ func TestZeroDowntimeUpgradeAfterOperatorUpgrade(t *testing.T) {
 	g.Expect(endpoints.Subsets[0].Addresses).To(HaveLen(1))
 
 	// Upgrade KubeRay operator to latest version and replace CRDs
-	test.T().Logf("Upgrading the KubeRay operator to the latest release")
-	cmd := exec.Command("kubectl", "replace", "-k", fmt.Sprintf("github.com/ray-project/kuberay/ray-operator/config/crd?ref=%s", upgradeVersion)) //nolint:gosec // required for upgrade
-	err = cmd.Run()
-	g.Expect(err).NotTo(HaveOccurred())
-	g.Eventually(cmd, TestTimeoutShort).Should(WithTransform(ProcessStateSuccess, BeTrue()))
-	cmd = exec.Command("helm", "upgrade", "kuberay-operator", "kuberay/kuberay-operator", "--version", upgradeVersion)
-	err = cmd.Run()
-	g.Expect(err).NotTo(HaveOccurred())
-	g.Eventually(cmd, TestTimeoutShort).Should(WithTransform(ProcessStateSuccess, BeTrue()))
+	test.T().Logf("Upgrading the KubeRay operator to %s", upgradeVersion)
+
+	useLocalHelmChart := false
+
+	if val, ok := os.LookupEnv("USE_LOCAL_HELM_CHART"); ok && val == "true" {
+		useLocalHelmChart = true
+	}
+
+	if useLocalHelmChart {
+		cmd := exec.Command("kubectl", "replace", "-k", "../../config/crd")
+		output, err := cmd.CombinedOutput()
+		if err != nil {
+			test.T().Logf("Error upgrading KubeRay operator: %v\nCommand output:\n%s", err, string(output))
+			t.Fatal("Failed to upgrade KubeRay operator")
+		}
+		g.Eventually(cmd, TestTimeoutShort).Should(WithTransform(ProcessStateSuccess, BeTrue()))
+
+		t.Logf("Upgrading operator deployment using local helm image version: %s", upgradeVersion)
+		cmd = exec.Command("helm", "upgrade", "kuberay-operator", "../../../helm-chart/kuberay-operator", "--set", fmt.Sprintf("image.repository=kuberay/kuberay-operator,image.tag=%s", upgradeVersion)) //nolint:gosec // required for upgrade
+		err = cmd.Run()
+		g.Expect(err).NotTo(HaveOccurred())
+		g.Eventually(cmd, TestTimeoutShort).Should(WithTransform(ProcessStateSuccess, BeTrue()))
+	} else {
+		cmd := exec.Command("kubectl", "replace", "-k", fmt.Sprintf("github.com/ray-project/kuberay/ray-operator/config/crd?ref=%s", upgradeVersion)) //nolint:gosec // required for upgrade
+		err = cmd.Run()
+		g.Expect(err).NotTo(HaveOccurred())
+		g.Eventually(cmd, TestTimeoutShort).Should(WithTransform(ProcessStateSuccess, BeTrue()))
+		cmd = exec.Command("helm", "upgrade", "kuberay-operator", "kuberay/kuberay-operator", "--version", upgradeVersion)
+		err = cmd.Run()
+		g.Expect(err).NotTo(HaveOccurred())
+		g.Eventually(cmd, TestTimeoutShort).Should(WithTransform(ProcessStateSuccess, BeTrue()))
+	}
 
 	// Validate RayService is able to serve requests during the upgrade
 	test.T().Logf("Sending requests to the RayService to make sure it is ready to serve requests")