add support for image override for argocd-extension

common/keys.go

@@ -172,6 +172,10 @@ const (
 	// to used for the argocd container.
 	ArgoCDImageEnvName = "ARGOCD_IMAGE"
 
+	// ArgoCDExtensionImageEnvName is the environment variable used to get the image
+	// to be used for the Argo-CD extension image
+	ArgoCDExtensionImageEnvName = "ARGOCD_EXTENSION_IMAGE"
+
 	// ArgoCDKeycloakImageEnvName is the environment variable used to get the image
 	// to used for the Keycloak container.
 	ArgoCDKeycloakImageEnvName = "ARGOCD_KEYCLOAK_IMAGE"

controllers/argocd/deployment.go

@@ -1564,16 +1564,9 @@ func updateNodePlacement(existing *appsv1.Deployment, deploy *appsv1.Deployment,
 }
 
 func getRolloutInitContainer() []corev1.Container {
-	return []corev1.Container{
+	containers := []corev1.Container{
 		{
-			Name:  "rollout-extension",
-			Image: common.ArgoCDExtensionInstallerImage,
-			Env: []corev1.EnvVar{
-				{
-					Name:  "EXTENSION_URL",
-					Value: common.ArgoRolloutsExtensionURL,
-				},
-			},
+			Name: "rollout-extension",
 			VolumeMounts: []corev1.VolumeMount{
 				{
 					Name:      "extensions",
@@ -1593,6 +1586,18 @@ func getRolloutInitContainer() []corev1.Container {
 			},
 		},
 	}
+
+	if value, exists := os.LookupEnv(common.ArgoCDExtensionImageEnvName); exists {
+		containers[0].Image = value
+	} else {
+		containers[0].Image = common.ArgoCDExtensionInstallerImage
+		containers[0].Env = []corev1.EnvVar{
+			{
+				Name:  "EXTENSION_URL",
+				Value: common.ArgoRolloutsExtensionURL,
+			}}
+	}
+	return containers
 }
 
 func removeInitContainer(initContainers []corev1.Container, name string) []corev1.Container {

controllers/argocd/deployment_test.go

@@ -3,6 +3,7 @@ package argocd
 import (
 	"context"
 	"fmt"
+	"os"
 	"reflect"
 	"strings"
 	"testing"
@@ -2409,3 +2410,44 @@ func TestReconcileArgoCD_reconcileRepoServerWithRemote(t *testing.T) {
 	assert.NoError(t, r.reconcileRepoDeployment(cr, false))
 	assert.NoError(t, r.Client.Get(context.TODO(), types.NamespacedName{Name: cr.Name + "-repo-server", Namespace: cr.Namespace}, d))
 }
+
+func Test_getRolloutInitContainer(t *testing.T) {
+	tests := []struct {
+		name      string
+		envSet    bool
+		wantImage string
+		wantEnv   []corev1.EnvVar
+	}{
+		{
+			name:      "when running in argocd-operator",
+			envSet:    false,
+			wantImage: "quay.io/argoprojlabs/argocd-extension-installer:v0.0.8",
+			wantEnv: []corev1.EnvVar{
+				{
+					Name:  "EXTENSION_URL",
+					Value: common.ArgoRolloutsExtensionURL,
+				},
+			},
+		},
+		{
+			name:      "when running in gitops-operator",
+			envSet:    true,
+			wantImage: "updated_container",
+			wantEnv:   nil,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.envSet {
+				err := os.Setenv(common.ArgoCDExtensionImageEnvName, "updated_container")
+				assert.NoError(t, err)
+			}
+
+			containers := getRolloutInitContainer()
+
+			assert.Equalf(t, tt.wantImage, containers[0].Image, "Image check")
+			assert.Equalf(t, tt.wantEnv, containers[0].Env, "Env check")
+
+		})
+	}
+}