From 866baec66e68c8904571be27c8f25dfca6b3dc79 Mon Sep 17 00:00:00 2001
From: Ross Peoples <rpeoples@redhat.com>
Date: Thu, 16 Dec 2021 14:06:46 -0600
Subject: [PATCH] Allow custom img resolver and flags for set image

Kubernetes-commit: 2f56584853355c0ab16e5ad2b0195d8dfcd7dc6b
---
 pkg/cmd/set/set_image.go      | 18 +++++++++++-------
 pkg/cmd/set/set_image_test.go | 15 +++++++++++++++
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/pkg/cmd/set/set_image.go b/pkg/cmd/set/set_image.go
index f25f6926..2df98035 100644
--- a/pkg/cmd/set/set_image.go
+++ b/pkg/cmd/set/set_image.go
@@ -51,7 +51,7 @@ type SetImageOptions struct {
 	All            bool
 	Output         string
 	Local          bool
-	ResolveImage   ImageResolver
+	ResolveImage   ImageResolverFunc
 	fieldManager   string
 
 	PrintObj printers.ResourcePrinterFunc
@@ -64,6 +64,14 @@ type SetImageOptions struct {
 	genericclioptions.IOStreams
 }
 
+// ImageResolver is a func that receives an image name, and
+// resolves it to an appropriate / compatible image name.
+// Adds flexibility for future image resolving methods.
+type ImageResolverFunc func(in string) (string, error)
+
+// ImageResolver to use.
+var ImageResolver = resolveImageFunc
+
 var (
 	imageResources = i18n.T(`
   	pod (po), replicationcontroller (rc), deployment (deploy), daemonset (ds), statefulset (sts), cronjob (cj), replicaset (rs)`)
@@ -127,6 +135,7 @@ func NewCmdImage(f cmdutil.Factory, streams genericclioptions.IOStreams) *cobra.
 	cmd.Flags().BoolVar(&o.Local, "local", o.Local, "If true, set image will NOT contact api-server but run locally.")
 	cmdutil.AddDryRunFlag(cmd)
 	cmdutil.AddFieldManagerFlagVar(cmd, &o.fieldManager, "kubectl-set")
+
 	return cmd
 }
 
@@ -151,7 +160,7 @@ func (o *SetImageOptions) Complete(f cmdutil.Factory, cmd *cobra.Command, args [
 	}
 	o.DryRunVerifier = resource.NewDryRunVerifier(dynamicClient, f.OpenAPIGetter())
 	o.Output = cmdutil.GetFlagString(cmd, "output")
-	o.ResolveImage = resolveImageFunc
+	o.ResolveImage = ImageResolver
 
 	cmdutil.PrintFlagsWithDryRunStrategy(o.PrintFlags, o.DryRunStrategy)
 	printer, err := o.PrintFlags.ToPrinter()
@@ -326,11 +335,6 @@ func hasWildcardKey(containerImages map[string]string) bool {
 	return ok
 }
 
-// ImageResolver is a func that receives an image name, and
-// resolves it to an appropriate / compatible image name.
-// Adds flexibility for future image resolving methods.
-type ImageResolver func(in string) (string, error)
-
 // implements ImageResolver
 func resolveImageFunc(in string) (string, error) {
 	return in, nil
diff --git a/pkg/cmd/set/set_image_test.go b/pkg/cmd/set/set_image_test.go
index 6ebe623e..28915791 100644
--- a/pkg/cmd/set/set_image_test.go
+++ b/pkg/cmd/set/set_image_test.go
@@ -778,3 +778,18 @@ func TestSetImageRemoteWithSpecificContainers(t *testing.T) {
 		})
 	}
 }
+
+func TestSetImageResolver(t *testing.T) {
+	f := func(in string) (string, error) {
+		return "custom", nil
+	}
+
+	ImageResolver = f
+
+	out, err := ImageResolver("my-image")
+	if err != nil {
+		t.Errorf("unexpected error from ImageResolver: %v", err)
+	} else if out != "custom" {
+		t.Errorf("expected: %s, found: %s", "custom", out)
+	}
+}