From 9df847b7ed6df2bcde222b407e4fbca4b96dfc58 Mon Sep 17 00:00:00 2001
From: Angelos Kolaitis <angelos.kolaitis@canonical.com>
Date: Sun, 2 Jun 2024 22:13:53 +0300
Subject: [PATCH] configurable metrics-server image through private annotations

---
 .../k8sd/features/metrics-server/internal.go  | 29 +++++++++++++++++++
 .../features/metrics-server/metrics_server.go |  6 ++--
 .../metrics-server/metrics_server_test.go     | 27 ++++++++++++++++-
 3 files changed, 59 insertions(+), 3 deletions(-)
 create mode 100644 src/k8s/pkg/k8sd/features/metrics-server/internal.go

diff --git a/src/k8s/pkg/k8sd/features/metrics-server/internal.go b/src/k8s/pkg/k8sd/features/metrics-server/internal.go
new file mode 100644
index 000000000..2300b7139
--- /dev/null
+++ b/src/k8s/pkg/k8sd/features/metrics-server/internal.go
@@ -0,0 +1,29 @@
+package metrics_server
+
+import "github.com/canonical/k8s/pkg/k8sd/types"
+
+const (
+	annotationImageRepo = "k8sd/v1alpha1/metrics-server/image-repo"
+	annotationImageTag  = "k8sd/v1alpha1/metrics-server/image-tag"
+)
+
+type config struct {
+	imageRepo string
+	imageTag  string
+}
+
+func internalConfig(annotations types.Annotations) config {
+	config := config{
+		imageRepo: imageRepo,
+		imageTag:  imageTag,
+	}
+
+	if v, ok := annotations.Get(annotationImageRepo); ok {
+		config.imageRepo = v
+	}
+	if v, ok := annotations.Get(annotationImageTag); ok {
+		config.imageTag = v
+	}
+
+	return config
+}
diff --git a/src/k8s/pkg/k8sd/features/metrics-server/metrics_server.go b/src/k8s/pkg/k8sd/features/metrics-server/metrics_server.go
index 873c44f8a..224db024e 100644
--- a/src/k8s/pkg/k8sd/features/metrics-server/metrics_server.go
+++ b/src/k8s/pkg/k8sd/features/metrics-server/metrics_server.go
@@ -14,10 +14,12 @@ import (
 func ApplyMetricsServer(ctx context.Context, snap snap.Snap, cfg types.MetricsServer, annotations types.Annotations) error {
 	m := snap.HelmClient()
 
+	config := internalConfig(annotations)
+
 	values := map[string]any{
 		"image": map[string]any{
-			"repository": imageRepo,
-			"tag":        imageTag,
+			"repository": config.imageRepo,
+			"tag":        config.imageTag,
 		},
 		"securityContext": map[string]any{
 			// ROCKs with Pebble as the entrypoint do not work with this option.
diff --git a/src/k8s/pkg/k8sd/features/metrics-server/metrics_server_test.go b/src/k8s/pkg/k8sd/features/metrics-server/metrics_server_test.go
index 855bdd7bb..c537211c4 100644
--- a/src/k8s/pkg/k8sd/features/metrics-server/metrics_server_test.go
+++ b/src/k8s/pkg/k8sd/features/metrics-server/metrics_server_test.go
@@ -43,7 +43,7 @@ func TestApplyMetricsServer(t *testing.T) {
 				},
 			}
 
-			err := metrics_server.ApplyMetricsServer(context.Background(), s, tc.config)
+			err := metrics_server.ApplyMetricsServer(context.Background(), s, tc.config, nil)
 			g.Expect(err).ToNot(HaveOccurred())
 
 			g.Expect(h.ApplyCalledWith).To(ConsistOf(SatisfyAll(
@@ -53,4 +53,29 @@ func TestApplyMetricsServer(t *testing.T) {
 			)))
 		})
 	}
+
+	t.Run("Annotations", func(t *testing.T) {
+		g := NewWithT(t)
+		h := &helmmock.Mock{}
+		s := &snapmock.Snap{
+			Mock: snapmock.Mock{
+				HelmClient: h,
+			},
+		}
+
+		cfg := types.MetricsServer{
+			Enabled: utils.Pointer(true),
+		}
+		annotations := types.Annotations{
+			"k8sd/v1alpha1/metrics-server/image-repo": "custom-image",
+			"k8sd/v1alpha1/metrics-server/image-tag":  "custom-tag",
+		}
+
+		err := metrics_server.ApplyMetricsServer(context.Background(), s, cfg, annotations)
+		g.Expect(err).To(BeNil())
+		g.Expect(h.ApplyCalledWith).To(ConsistOf(HaveField("Values", HaveKeyWithValue("image", SatisfyAll(
+			HaveKeyWithValue("repository", "custom-image"),
+			HaveKeyWithValue("tag", "custom-tag"),
+		)))))
+	})
 }