From 2bd2a0ab53eed258af11beee1cf8c2537a3995dc Mon Sep 17 00:00:00 2001
From: Moritz Clasmeier <mtesseract@users.noreply.github.com>
Date: Thu, 27 Jul 2023 18:40:39 +0200
Subject: [PATCH] Expose generic mechanism for configuring the controller,
 allowing to add additional watch sources. (#153)

Expose generic mechanism for configuring the controller, allowing to add additional watch sources.
Signed-off-by: Moritz Clasmeier <moritz@stackrox.com>
---
 pkg/reconciler/reconciler.go      | 32 +++++++++++++++++++++++++
 pkg/reconciler/reconciler_test.go | 39 +++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+)

diff --git a/pkg/reconciler/reconciler.go b/pkg/reconciler/reconciler.go
index f8fc1286..097581ff 100644
--- a/pkg/reconciler/reconciler.go
+++ b/pkg/reconciler/reconciler.go
@@ -80,6 +80,7 @@ type Reconciler struct {
 	reconcilePeriod                  time.Duration
 	maxHistory                       int
 	skipPrimaryGVKSchemeRegistration bool
+	controllerSetupFuncs             []ControllerSetupFunc
 
 	annotSetupOnce       sync.Once
 	annotations          map[string]struct{}
@@ -149,6 +150,13 @@ func (r *Reconciler) SetupWithManager(mgr ctrl.Manager) error {
 		return err
 	}
 
+	for _, f := range r.controllerSetupFuncs {
+		err = f(c)
+		if err != nil {
+			return fmt.Errorf("failed to execute custom controller setup function: %v", err)
+		}
+	}
+
 	r.log.Info("Watching resource",
 		"group", r.gvk.Group,
 		"version", r.gvk.Version,
@@ -482,6 +490,30 @@ func WithSelector(s metav1.LabelSelector) Option {
 	}
 }
 
+// WithControllerSetupFunc is an Option that allows customizing a controller before it is started.
+// The only supported customization here is adding additional Watch sources to the controller.
+func WithControllerSetupFunc(f ControllerSetupFunc) Option {
+	return func(r *Reconciler) error {
+		r.controllerSetupFuncs = append(r.controllerSetupFuncs, f)
+		return nil
+	}
+}
+
+// ControllerSetup allows restricted access to the Controller using the WithControllerSetupFunc option.
+// Currently the only supposed configuration is adding additional watchers do the controller.
+type ControllerSetup interface {
+	// Watch takes events provided by a Source and uses the EventHandler to
+	// enqueue reconcile.Requests in response to the events.
+	//
+	// Watch may be provided one or more Predicates to filter events before
+	// they are given to the EventHandler.  Events will be passed to the
+	// EventHandler if all provided Predicates evaluate to true.
+	Watch(src source.Source, eventhandler handler.EventHandler, predicates ...predicate.Predicate) error
+}
+
+// ControllerSetupFunc allows configuring a controller's builder.
+type ControllerSetupFunc func(c ControllerSetup) error
+
 // Reconcile reconciles a CR that defines a Helm v3 release.
 //
 //   - If a release does not exist for this CR, a new release is installed.
diff --git a/pkg/reconciler/reconciler_test.go b/pkg/reconciler/reconciler_test.go
index 2ef4d98e..323df098 100644
--- a/pkg/reconciler/reconciler_test.go
+++ b/pkg/reconciler/reconciler_test.go
@@ -48,6 +48,7 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/log/zap"
 	"sigs.k8s.io/controller-runtime/pkg/manager"
 	"sigs.k8s.io/controller-runtime/pkg/reconcile"
+	"sigs.k8s.io/controller-runtime/pkg/source"
 	"sigs.k8s.io/yaml"
 
 	"github.com/operator-framework/helm-operator-plugins/internal/sdk/controllerutil"
@@ -59,6 +60,7 @@ import (
 	"github.com/operator-framework/helm-operator-plugins/pkg/reconciler/internal/conditions"
 	helmfake "github.com/operator-framework/helm-operator-plugins/pkg/reconciler/internal/fake"
 	"github.com/operator-framework/helm-operator-plugins/pkg/values"
+	sdkhandler "github.com/operator-framework/operator-lib/handler"
 )
 
 // custom is used within the reconciler test suite as underlying type for the GVK scheme.
@@ -1328,6 +1330,43 @@ var _ = Describe("Reconciler", func() {
 		})
 
 	})
+
+	var _ = Describe("Test custom controller setup", func() {
+		var (
+			mgr                   manager.Manager
+			r                     *Reconciler
+			err                   error
+			controllerSetupCalled bool
+		)
+		additionalGVK := schema.GroupVersionKind{Group: "example.com", Version: "v1", Kind: "SomeOtherKind"}
+		setupController := func(c ControllerSetup) error {
+			controllerSetupCalled = true
+			u := &unstructured.Unstructured{}
+			u.SetGroupVersionKind(additionalGVK)
+			return c.Watch(&source.Kind{Type: u}, &sdkhandler.InstrumentedEnqueueRequestForObject{})
+		}
+
+		It("Registering builder setup function for reconciler works", func() {
+			mgr = getManagerOrFail()
+			r, err = New(
+				WithGroupVersionKind(gvk),
+				WithChart(chrt),
+				WithInstallAnnotations(annotation.InstallDescription{}),
+				WithUpgradeAnnotations(annotation.UpgradeDescription{}),
+				WithUninstallAnnotations(annotation.UninstallDescription{}),
+				WithOverrideValues(map[string]string{
+					"image.repository": "custom-nginx",
+				}),
+				WithControllerSetupFunc(setupController),
+			)
+			Expect(err).To(BeNil())
+		})
+
+		It("Setting up reconciler with manager causes custom builder setup to be executed", func() {
+			Expect(r.SetupWithManager(mgr)).To(Succeed())
+			Expect(controllerSetupCalled).To(BeTrue())
+		})
+	})
 })
 
 func getManagerOrFail() manager.Manager {