From 70f7fd6b000be3b321a840a97b35731ef791a9b8 Mon Sep 17 00:00:00 2001
From: Kimmo Lehto <kimmo.lehto@gmail.com>
Date: Wed, 5 Aug 2020 13:22:33 +0300
Subject: [PATCH] Repo-agnostic retag in PullImages (#154)

* Repo-agnostic retag in PullImages

* Lints
---
 pkg/phase/pull_images.go      |  6 +++++-
 pkg/phase/pull_images_test.go | 12 ++++++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)
 create mode 100644 pkg/phase/pull_images_test.go

diff --git a/pkg/phase/pull_images.go b/pkg/phase/pull_images.go
index 15bcbe508..bb9c1c82e 100644
--- a/pkg/phase/pull_images.go
+++ b/pkg/phase/pull_images.go
@@ -39,7 +39,7 @@ func (p *PullImages) Run(c *api.ClusterConfig) error {
 		// Store map of original --> custom repo images
 		imageMap := make(map[string]string, len(images))
 		for index, i := range images {
-			newImage := strings.Replace(i, "docker/", fmt.Sprintf("%s/", c.Spec.Ucp.ImageRepo), 1)
+			newImage := p.ImageFromCustomRepo(i, c.Spec.Ucp.ImageRepo)
 			imageMap[i] = newImage
 			images[index] = newImage
 		}
@@ -56,7 +56,11 @@ func (p *PullImages) Run(c *api.ClusterConfig) error {
 	return runParallelOnHosts(c.Spec.Managers(), c, func(h *api.Host, c *api.ClusterConfig) error {
 		return p.pullImages(h, images)
 	})
+}
 
+// ImageFromCustomRepo will replace the organization part in an image name
+func (p *PullImages) ImageFromCustomRepo(image, repo string) string {
+	return fmt.Sprintf("%s%s", repo, image[strings.IndexByte(image, '/'):])
 }
 
 func (p *PullImages) listImages(config *api.ClusterConfig) ([]string, error) {
diff --git a/pkg/phase/pull_images_test.go b/pkg/phase/pull_images_test.go
new file mode 100644
index 000000000..6bf330d16
--- /dev/null
+++ b/pkg/phase/pull_images_test.go
@@ -0,0 +1,12 @@
+package phase
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestPullImagesCustomRepo(t *testing.T) {
+	phase := PullImages{}
+	require.Equal(t, "xyz/foofoo:1.2.3-latest", phase.ImageFromCustomRepo("docker/foofoo:1.2.3-latest", "xyz"))
+}