From a714810ec565e6bc3c0276426e27cdf2f5efe996 Mon Sep 17 00:00:00 2001
From: Michael Burman <yak@iki.fi>
Date: Fri, 11 Oct 2024 10:39:53 +0300
Subject: [PATCH] Add HCD version override test

---
 apis/config/v1beta1/imageconfig_types.go     |  2 ++
 apis/config/v1beta1/zz_generated.deepcopy.go |  7 +++++++
 pkg/images/images.go                         |  5 +++++
 pkg/images/images_test.go                    | 20 +++++++++++++++++---
 4 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/apis/config/v1beta1/imageconfig_types.go b/apis/config/v1beta1/imageconfig_types.go
index f06afda9..14b3a368 100644
--- a/apis/config/v1beta1/imageconfig_types.go
+++ b/apis/config/v1beta1/imageconfig_types.go
@@ -57,6 +57,8 @@ type Images struct {
 
 	DSEVersions map[string]string `json:"dse,omitempty"`
 
+	HCDVersions map[string]string `json:"hcd,omitempty"`
+
 	SystemLogger string `json:"system-logger,omitempty"`
 
 	Client string `json:"k8ssandra-client,omitempty"`
diff --git a/apis/config/v1beta1/zz_generated.deepcopy.go b/apis/config/v1beta1/zz_generated.deepcopy.go
index e7903da9..8559fadc 100644
--- a/apis/config/v1beta1/zz_generated.deepcopy.go
+++ b/apis/config/v1beta1/zz_generated.deepcopy.go
@@ -157,6 +157,13 @@ func (in *Images) DeepCopyInto(out *Images) {
 			(*out)[key] = val
 		}
 	}
+	if in.HCDVersions != nil {
+		in, out := &in.HCDVersions, &out.HCDVersions
+		*out = make(map[string]string, len(*in))
+		for key, val := range *in {
+			(*out)[key] = val
+		}
+	}
 	if in.Others != nil {
 		in, out := &in.Others, &out.Others
 		*out = make(map[string]string, len(*in))
diff --git a/pkg/images/images.go b/pkg/images/images.go
index bd575589..07bca5fa 100644
--- a/pkg/images/images.go
+++ b/pkg/images/images.go
@@ -173,6 +173,11 @@ func getCassandraContainerImageOverride(serverType, version string) (bool, strin
 				return true, value
 			}
 		}
+		if serverType == "hcd" {
+			if value, found := images.HCDVersions[version]; found {
+				return true, value
+			}
+		}
 	}
 	return false, ""
 }
diff --git a/pkg/images/images_test.go b/pkg/images/images_test.go
index cbc473a8..04fbb452 100644
--- a/pkg/images/images_test.go
+++ b/pkg/images/images_test.go
@@ -61,14 +61,28 @@ func TestCassandraOverride(t *testing.T) {
 	assert.NoError(err, "getting Cassandra image with overrides should succeed")
 	assert.Equal(fmt.Sprintf("ghcr.io/%s", customImageName), cassImage)
 
-	customImageWithOrg := "k8ssandra/cass-management-api:4.0.0"
+	customImageNamespace := "modified"
 	imageConfig.Images.CassandraVersions = map[string]string{
-		"4.0.0": fmt.Sprintf("us-docker.pkg.dev/%s", customImageWithOrg),
+		"4.0.0": fmt.Sprintf("us-docker.pkg.dev/%s/cass-management-api:4.0.0", customImageNamespace),
+	}
+	imageConfig.Images.DSEVersions = map[string]string{
+		"6.8.0": fmt.Sprintf("us-docker.pkg.dev/%s/dse-mgmtapi-6_8:6.8.0", customImageNamespace),
+	}
+	imageConfig.Images.HCDVersions = map[string]string{
+		"1.0.0": fmt.Sprintf("us-docker.pkg.dev/%s/hcd:1.0.0", customImageNamespace),
 	}
 
 	cassImage, err = GetCassandraImage("cassandra", "4.0.0")
 	assert.NoError(err, "getting Cassandra image with overrides should succeed")
-	assert.Equal(fmt.Sprintf("ghcr.io/%s", customImageWithOrg), cassImage)
+	assert.Equal("ghcr.io/modified/cass-management-api:4.0.0", cassImage)
+
+	cassImage, err = GetCassandraImage("dse", "6.8.0")
+	assert.NoError(err, "getting Cassandra image with overrides should succeed")
+	assert.Equal("ghcr.io/modified/dse-mgmtapi-6_8:6.8.0", cassImage)
+
+	cassImage, err = GetCassandraImage("hcd", "1.0.0")
+	assert.NoError(err, "getting Cassandra image with overrides should succeed")
+	assert.Equal("ghcr.io/modified/hcd:1.0.0", cassImage)
 }
 
 func TestDefaultImageConfigParsing(t *testing.T) {