Add various functional options to ActionConfigGetter and ActionClientGetter constructors

pkg/client/actionclient.go

@@ -61,13 +61,71 @@ type GetOption func(*action.Get) error
 type InstallOption func(*action.Install) error
 type UpgradeOption func(*action.Upgrade) error
 type UninstallOption func(*action.Uninstall) error
+type RollbackOption func(*action.Rollback) error
 
-func NewActionClientGetter(acg ActionConfigGetter) ActionClientGetter {
-	return &actionClientGetter{acg}
+type ActionClientGetterOption func(*actionClientGetter) error
+
+func AppendGetOptions(opts ...GetOption) ActionClientGetterOption {
+	return func(getter *actionClientGetter) error {
+		getter.defaultGetOpts = append(getter.defaultGetOpts, opts...)
+		return nil
+	}
+}
+
+func AppendInstallOptions(opts ...InstallOption) ActionClientGetterOption {
+	return func(getter *actionClientGetter) error {
+		getter.defaultInstallOpts = append(getter.defaultInstallOpts, opts...)
+		return nil
+	}
+}
+
+func AppendUpgradeOptions(opts ...UpgradeOption) ActionClientGetterOption {
+	return func(getter *actionClientGetter) error {
+		getter.defaultUpgradeOpts = append(getter.defaultUpgradeOpts, opts...)
+		return nil
+	}
+}
+
+func AppendUninstallOptions(opts ...UninstallOption) ActionClientGetterOption {
+	return func(getter *actionClientGetter) error {
+		getter.defaultUninstallOpts = append(getter.defaultUninstallOpts, opts...)
+		return nil
+	}
+}
+
+func AppendInstallFailureUninstallOptions(opts ...UninstallOption) ActionClientGetterOption {
+	return func(getter *actionClientGetter) error {
+		getter.installFailureUninstallOpts = append(getter.installFailureUninstallOpts, opts...)
+		return nil
+	}
+}
+func AppendUpgradeFailureRollbackOptions(opts ...RollbackOption) ActionClientGetterOption {
+	return func(getter *actionClientGetter) error {
+		getter.upgradeFailureRollbackOpts = append(getter.upgradeFailureRollbackOpts, opts...)
+		return nil
+	}
+}
+
+func NewActionClientGetter(acg ActionConfigGetter, opts ...ActionClientGetterOption) (ActionClientGetter, error) {
+	actionClientGetter := &actionClientGetter{acg: acg}
+	for _, opt := range opts {
+		if err := opt(actionClientGetter); err != nil {
+			return nil, err
+		}
+	}
+	return actionClientGetter, nil
 }
 
 type actionClientGetter struct {
 	acg ActionConfigGetter
+
+	defaultGetOpts       []GetOption
+	defaultInstallOpts   []InstallOption
+	defaultUpgradeOpts   []UpgradeOption
+	defaultUninstallOpts []UninstallOption
+
+	installFailureUninstallOpts []UninstallOption
+	upgradeFailureRollbackOpts  []RollbackOption
 }
 
 var _ ActionClientGetter = &actionClientGetter{}
@@ -83,9 +141,18 @@ func (hcg *actionClientGetter) ActionClientFor(obj client.Object) (ActionInterfa
 	}
 	postRenderer := DefaultPostRendererFunc(rm, actionConfig.KubeClient, obj)
 	return &actionClient{
-		conf:               actionConfig,
-		defaultInstallOpts: []InstallOption{WithInstallPostRenderer(postRenderer)},
-		defaultUpgradeOpts: []UpgradeOption{WithUpgradePostRenderer(postRenderer)},
+		conf: actionConfig,
+
+		// For the install and upgrade options, we put the post renderer first in the list
+		// on purpose because we want user-provided defaults to be able to override the
+		// post-renderer that we automatically configure for the client.
+		defaultGetOpts:       hcg.defaultGetOpts,
+		defaultInstallOpts:   append([]InstallOption{WithInstallPostRenderer(postRenderer)}, hcg.defaultInstallOpts...),
+		defaultUpgradeOpts:   append([]UpgradeOption{WithUpgradePostRenderer(postRenderer)}, hcg.defaultUpgradeOpts...),
+		defaultUninstallOpts: hcg.defaultUninstallOpts,
+
+		installFailureUninstallOpts: hcg.installFailureUninstallOpts,
+		upgradeFailureRollbackOpts:  hcg.upgradeFailureRollbackOpts,
 	}, nil
 }
 
@@ -96,6 +163,9 @@ type actionClient struct {
 	defaultInstallOpts   []InstallOption
 	defaultUpgradeOpts   []UpgradeOption
 	defaultUninstallOpts []UninstallOption
+
+	installFailureUninstallOpts []UninstallOption
+	upgradeFailureRollbackOpts  []RollbackOption
 }
 
 var _ ActionInterface = &actionClient{}
@@ -137,7 +207,7 @@ func (c *actionClient) Install(name, namespace string, chrt *chart.Chart, vals m
 			//
 			// Only return an error about a rollback failure if the failure was
 			// caused by something other than the release not being found.
-			_, uninstallErr := c.Uninstall(name)
+			_, uninstallErr := c.uninstall(name, c.installFailureUninstallOpts...)
 			if uninstallErr != nil && !errors.Is(uninstallErr, driver.ErrReleaseNotFound) {
 				return nil, fmt.Errorf("uninstall failed: %v: original install error: %w", uninstallErr, err)
 			}
@@ -158,16 +228,18 @@ func (c *actionClient) Upgrade(name, namespace string, chrt *chart.Chart, vals m
 	rel, err := upgrade.Run(name, chrt, vals)
 	if err != nil {
 		if rel != nil {
-			rollback := action.NewRollback(c.conf)
-			rollback.Force = true
-			rollback.MaxHistory = upgrade.MaxHistory
+			rollbackOpts := append([]RollbackOption{func(rollback *action.Rollback) error {
+				rollback.Force = true
+				rollback.MaxHistory = upgrade.MaxHistory
+				return nil
+			}}, c.upgradeFailureRollbackOpts...)
 
 			// As of Helm 2.13, if Upgrade returns a non-nil release, that
 			// means the release was also recorded in the release store.
 			// Therefore, we should perform the rollback when we have a non-nil
 			// release. Any rollback error here would be unexpected, so always
 			// log both the update and rollback errors.
-			rollbackErr := rollback.Run(name)
+			rollbackErr := c.rollback(name, rollbackOpts...)
 			if rollbackErr != nil {
 				return nil, fmt.Errorf("rollback failed: %v: original upgrade error: %w", rollbackErr, err)
 			}
@@ -177,9 +249,23 @@ func (c *actionClient) Upgrade(name, namespace string, chrt *chart.Chart, vals m
 	return rel, nil
 }
 
+func (c *actionClient) rollback(name string, opts ...RollbackOption) error {
+	rollback := action.NewRollback(c.conf)
+	for _, o := range opts {
+		if err := o(rollback); err != nil {
+			return err
+		}
+	}
+	return rollback.Run(name)
+}
+
 func (c *actionClient) Uninstall(name string, opts ...UninstallOption) (*release.UninstallReleaseResponse, error) {
+	return c.uninstall(name, concat(c.defaultUninstallOpts, opts...)...)
+}
+
+func (c *actionClient) uninstall(name string, opts ...UninstallOption) (*release.UninstallReleaseResponse, error) {
 	uninstall := action.NewUninstall(c.conf)
-	for _, o := range concat(c.defaultUninstallOpts, opts...) {
+	for _, o := range opts {
 		if err := o(uninstall); err != nil {
 			return nil, err
 		}

pkg/client/actionclient_test.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"errors"
 	"strconv"
+	"time"
 
 	"github.com/go-logr/logr"
 	. "github.com/onsi/ginkgo/v2"
@@ -37,6 +38,7 @@ import (
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	apitypes "k8s.io/apimachinery/pkg/types"
+	"k8s.io/apimachinery/pkg/util/rand"
 	"k8s.io/cli-runtime/pkg/resource"
 	"k8s.io/utils/pointer"
 	"sigs.k8s.io/controller-runtime/pkg/client"
@@ -61,7 +63,178 @@ var _ = Describe("ActionClient", func() {
 		It("should return a valid ActionConfigGetter", func() {
 			actionConfigGetter, err := NewActionConfigGetter(cfg, rm, logr.Discard())
 			Expect(err).ShouldNot(HaveOccurred())
-			Expect(NewActionClientGetter(actionConfigGetter)).NotTo(BeNil())
+			acg, err := NewActionClientGetter(actionConfigGetter)
+			Expect(err).To(BeNil())
+			Expect(acg).NotTo(BeNil())
+		})
+
+		When("options are specified", func() {
+			expectErr := errors.New("expect this error")
+
+			var (
+				actionConfigGetter ActionConfigGetter
+				obj                client.Object
+			)
+			BeforeEach(func() {
+				var err error
+				actionConfigGetter, err = NewActionConfigGetter(cfg, rm, logr.Discard())
+				Expect(err).ShouldNot(HaveOccurred())
+				obj = testutil.BuildTestCR(gvk)
+			})
+
+			It("should get clients with custom get options", func() {
+				expectVersion := rand.Int()
+				acg, err := NewActionClientGetter(actionConfigGetter, AppendGetOptions(
+					func(get *action.Get) error {
+						get.Version = expectVersion
+						return nil
+					},
+					func(get *action.Get) error {
+						Expect(get.Version).To(Equal(expectVersion))
+						return expectErr
+					},
+				))
+				Expect(err).To(BeNil())
+				Expect(acg).NotTo(BeNil())
+
+				ac, err := acg.ActionClientFor(obj)
+				Expect(err).To(BeNil())
+				Expect(ac).NotTo(BeNil())
+
+				_, err = ac.Get(obj.GetName())
+				Expect(err).To(MatchError(expectErr))
+			})
+			It("should get clients with custom install options", func() {
+				acg, err := NewActionClientGetter(actionConfigGetter, AppendInstallOptions(
+					func(install *action.Install) error {
+						install.Description = mockTestDesc
+						return nil
+					},
+					func(install *action.Install) error {
+						Expect(install.Description).To(Equal(mockTestDesc))
+						return expectErr
+					},
+				))
+				Expect(err).To(BeNil())
+				Expect(acg).NotTo(BeNil())
+
+				ac, err := acg.ActionClientFor(obj)
+				Expect(err).To(BeNil())
+				Expect(ac).NotTo(BeNil())
+
+				_, err = ac.Install(obj.GetName(), obj.GetNamespace(), &chrt, chartutil.Values{})
+				Expect(err).To(MatchError(expectErr))
+			})
+			It("should get clients with custom upgrade options", func() {
+				acg, err := NewActionClientGetter(actionConfigGetter, AppendUpgradeOptions(
+					func(upgrade *action.Upgrade) error {
+						upgrade.Description = mockTestDesc
+						return nil
+					},
+					func(upgrade *action.Upgrade) error {
+						Expect(upgrade.Description).To(Equal(mockTestDesc))
+						return expectErr
+					},
+				))
+				Expect(err).To(BeNil())
+				Expect(acg).NotTo(BeNil())
+
+				ac, err := acg.ActionClientFor(obj)
+				Expect(err).To(BeNil())
+				Expect(ac).NotTo(BeNil())
+
+				_, err = ac.Upgrade(obj.GetName(), obj.GetNamespace(), &chrt, chartutil.Values{})
+				Expect(err).To(MatchError(expectErr))
+			})
+			It("should get clients with custom uninstall options", func() {
+				acg, err := NewActionClientGetter(actionConfigGetter, AppendUninstallOptions(
+					func(uninstall *action.Uninstall) error {
+						uninstall.Description = mockTestDesc
+						return nil
+					},
+					func(uninstall *action.Uninstall) error {
+						Expect(uninstall.Description).To(Equal(mockTestDesc))
+						return expectErr
+					},
+				))
+				Expect(err).To(BeNil())
+				Expect(acg).NotTo(BeNil())
+
+				ac, err := acg.ActionClientFor(obj)
+				Expect(err).To(BeNil())
+				Expect(ac).NotTo(BeNil())
+
+				_, err = ac.Uninstall(obj.GetName())
+				Expect(err).To(MatchError(expectErr))
+			})
+			It("should get clients with custom install failure uninstall options", func() {
+				acg, err := NewActionClientGetter(actionConfigGetter, AppendInstallFailureUninstallOptions(
+					func(uninstall *action.Uninstall) error {
+						uninstall.Description = mockTestDesc
+						return nil
+					},
+					func(uninstall *action.Uninstall) error {
+						Expect(uninstall.Description).To(Equal(mockTestDesc))
+						return expectErr
+					},
+				))
+				Expect(err).To(BeNil())
+				Expect(acg).NotTo(BeNil())
+
+				ac, err := acg.ActionClientFor(obj)
+				Expect(err).To(BeNil())
+				Expect(ac).NotTo(BeNil())
+
+				_, err = ac.Install(obj.GetName(), obj.GetNamespace(), &chrt, chartutil.Values{}, func(install *action.Install) error {
+					// Force the installatiom to fail by using an impossibly short wait.
+					// When the installation fails, the failure uninstall logic is attempted.
+					install.Wait = true
+					install.Timeout = time.Nanosecond * 1
+					return nil
+				})
+				Expect(err).To(MatchError(ContainSubstring(expectErr.Error())))
+
+				// Uninstall the chart to cleanup for other tests.
+				_, err = ac.Uninstall(obj.GetName())
+				Expect(err).To(BeNil())
+			})
+			It("should get clients with custom upgrade failure rollback options", func() {
+				expectMaxHistory := rand.Int()
+				acg, err := NewActionClientGetter(actionConfigGetter, AppendUpgradeFailureRollbackOptions(
+					func(rollback *action.Rollback) error {
+						rollback.MaxHistory = expectMaxHistory
+						return nil
+					},
+					func(rollback *action.Rollback) error {
+						Expect(rollback.MaxHistory).To(Equal(expectMaxHistory))
+						return expectErr
+					},
+				))
+				Expect(err).To(BeNil())
+				Expect(acg).NotTo(BeNil())
+
+				ac, err := acg.ActionClientFor(obj)
+				Expect(err).To(BeNil())
+				Expect(ac).NotTo(BeNil())
+
+				// Install the chart so that we can try an upgrade.
+				rel, err := ac.Install(obj.GetName(), obj.GetNamespace(), &chrt, chartutil.Values{})
+				Expect(err).To(BeNil())
+				Expect(rel).NotTo(BeNil())
+
+				_, err = ac.Upgrade(obj.GetName(), obj.GetNamespace(), &chrt, chartutil.Values{}, func(upgrade *action.Upgrade) error {
+					// Force the upgrade to fail by using an impossibly short wait.
+					// When the upgrade fails, the rollback logic is attempted.
+					upgrade.Wait = true
+					upgrade.Timeout = time.Nanosecond * 1
+					return nil
+				})
+				Expect(err).To(MatchError(ContainSubstring(expectErr.Error())))
+
+				// Uninstall the chart to cleanup for other tests.
+				_, err = ac.Uninstall(obj.GetName())
+				Expect(err).To(BeNil())
+			})
 		})
 	})
 
@@ -88,7 +261,8 @@ var _ = Describe("ActionClient", func() {
 		It("should return a valid ActionClient", func() {
 			actionConfGetter, err := NewActionConfigGetter(cfg, rm, logr.Discard())
 			Expect(err).ShouldNot(HaveOccurred())
-			acg := NewActionClientGetter(actionConfGetter)
+			acg, err := NewActionClientGetter(actionConfGetter)
+			Expect(err).To(BeNil())
 			ac, err := acg.ActionClientFor(obj)
 			Expect(err).To(BeNil())
 			Expect(ac).NotTo(BeNil())
@@ -107,7 +281,8 @@ var _ = Describe("ActionClient", func() {
 
 			actionConfigGetter, err := NewActionConfigGetter(cfg, rm, logr.Discard())
 			Expect(err).ShouldNot(HaveOccurred())
-			acg := NewActionClientGetter(actionConfigGetter)
+			acg, err := NewActionClientGetter(actionConfigGetter)
+			Expect(err).To(BeNil())
 			ac, err = acg.ActionClientFor(obj)
 			Expect(err).To(BeNil())