From 932c3d724acf1c7910a340b76fc4bb428749516e Mon Sep 17 00:00:00 2001
From: Anand Kumar Singh <anandrkskd@gmail.com>
Date: Mon, 4 Nov 2024 14:20:45 +0530
Subject: [PATCH 1/2] add support for image override for argocd-extension

Signed-off-by: Anand Kumar Singh <anandrkskd@gmail.com>
---
 common/keys.go                        |  4 +++
 controllers/argocd/deployment.go      | 23 +++++++++------
 controllers/argocd/deployment_test.go | 42 +++++++++++++++++++++++++++
 3 files changed, 60 insertions(+), 9 deletions(-)

diff --git a/common/keys.go b/common/keys.go
index 274e9bb4d..b6f02a744 100644
--- a/common/keys.go
+++ b/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"
diff --git a/controllers/argocd/deployment.go b/controllers/argocd/deployment.go
index 9cba2bf71..baf91fca6 100644
--- a/controllers/argocd/deployment.go
+++ b/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 {
diff --git a/controllers/argocd/deployment_test.go b/controllers/argocd/deployment_test.go
index 3c62de11e..b3e804997 100644
--- a/controllers/argocd/deployment_test.go
+++ b/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 gitops",
+			envSet:    true,
+			wantImage: "updated_container",
+			wantEnv:   nil,
+		},
+		{
+			name:      "when running in gitops",
+			envSet:    false,
+			wantImage: "quay.io/argoprojlabs/argocd-extension-installer:v0.0.8",
+			wantEnv: []corev1.EnvVar{
+				{
+					Name:  "EXTENSION_URL",
+					Value: common.ArgoRolloutsExtensionURL,
+				},
+			},
+		},
+	}
+	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")
+
+		})
+	}
+}

From e69966eb70c79a7f5060ceb0d18efc1222d08423 Mon Sep 17 00:00:00 2001
From: Anand Kumar Singh <anandrkskd@gmail.com>
Date: Tue, 5 Nov 2024 11:12:07 +0530
Subject: [PATCH 2/2] fix unit test failure

Signed-off-by: Anand Kumar Singh <anandrkskd@gmail.com>
---
 controllers/argocd/deployment_test.go | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/controllers/argocd/deployment_test.go b/controllers/argocd/deployment_test.go
index b3e804997..83122d71f 100644
--- a/controllers/argocd/deployment_test.go
+++ b/controllers/argocd/deployment_test.go
@@ -2419,13 +2419,7 @@ func Test_getRolloutInitContainer(t *testing.T) {
 		wantEnv   []corev1.EnvVar
 	}{
 		{
-			name:      "when running in gitops",
-			envSet:    true,
-			wantImage: "updated_container",
-			wantEnv:   nil,
-		},
-		{
-			name:      "when running in gitops",
+			name:      "when running in argocd-operator",
 			envSet:    false,
 			wantImage: "quay.io/argoprojlabs/argocd-extension-installer:v0.0.8",
 			wantEnv: []corev1.EnvVar{
@@ -2435,6 +2429,12 @@ func Test_getRolloutInitContainer(t *testing.T) {
 				},
 			},
 		},
+		{
+			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) {