From 6612f620dfd44c84d0dc5522df28265e6723f3e9 Mon Sep 17 00:00:00 2001
From: David Grove <groved@us.ibm.com>
Date: Tue, 5 Mar 2024 13:25:45 -0500
Subject: [PATCH] factor out controller initialization logic

---
 cmd/main.go                              | 33 +++----------------
 internal/controller/appwrapper_config.go | 42 ++++++++++++++++++++++++
 2 files changed, 46 insertions(+), 29 deletions(-)

diff --git a/cmd/main.go b/cmd/main.go
index 1aec0b8..d747a03 100644
--- a/cmd/main.go
+++ b/cmd/main.go
@@ -34,7 +34,6 @@ import (
 	metricsserver "sigs.k8s.io/controller-runtime/pkg/metrics/server"
 	"sigs.k8s.io/controller-runtime/pkg/webhook"
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
-	"sigs.k8s.io/kueue/pkg/controller/jobframework"
 
 	workloadv1beta2 "github.com/project-codeflare/appwrapper/api/v1beta2"
 	"github.com/project-codeflare/appwrapper/internal/controller"
@@ -133,38 +132,14 @@ func main() {
 		os.Exit(1)
 	}
 
-	if err := controller.WorkloadReconciler(
-		mgr.GetClient(),
-		mgr.GetEventRecorderFor("kueue"),
-		jobframework.WithManageJobsWithoutQueueName(config.ManageJobsWithoutQueueName),
-	).SetupWithManager(mgr); err != nil {
-		setupLog.Error(err, "Unable to create controller", "controller", "Workload")
-		os.Exit(1)
-	}
-
-	if err = (&controller.AppWrapperReconciler{
-		Client: mgr.GetClient(),
-		Scheme: mgr.GetScheme(),
-		Config: &config,
-	}).SetupWithManager(mgr); err != nil {
-		setupLog.Error(err, "unable to create controller", "controller", "AppWrapper")
-		os.Exit(1)
-	}
-	if os.Getenv("ENABLE_WEBHOOKS") != "false" {
-		if err = (&controller.AppWrapperWebhook{
-			Config: &config,
-		}).SetupWebhookWithManager(mgr); err != nil {
-			setupLog.Error(err, "unable to create webhook", "webhook", "AppWrapper")
-			os.Exit(1)
-		}
-	}
-	//+kubebuilder:scaffold:builder
 	ctx := ctrl.SetupSignalHandler()
-	if err := jobframework.SetupWorkloadOwnerIndex(ctx, mgr.GetFieldIndexer(), controller.GVK); err != nil {
-		setupLog.Error(err, "Setting up indexes", "GVK", controller.GVK)
+	err = controller.SetupWithManager(ctx, mgr, &config)
+	if err != nil {
+		setupLog.Error(err, "unable to start appwrapper controllers")
 		os.Exit(1)
 	}
 
+	//+kubebuilder:scaffold:builder
 	if err := mgr.AddHealthzCheck("healthz", healthz.Ping); err != nil {
 		setupLog.Error(err, "unable to set up health check")
 		os.Exit(1)
diff --git a/internal/controller/appwrapper_config.go b/internal/controller/appwrapper_config.go
index c9ef320..e27a743 100644
--- a/internal/controller/appwrapper_config.go
+++ b/internal/controller/appwrapper_config.go
@@ -16,7 +16,49 @@ limitations under the License.
 
 package controller
 
+import (
+	"context"
+	"fmt"
+	"os"
+
+	ctrl "sigs.k8s.io/controller-runtime"
+	"sigs.k8s.io/kueue/pkg/controller/jobframework"
+)
+
 type AppWrapperConfig struct {
 	ManageJobsWithoutQueueName bool   `json:"manageJobsWithoutQueueName,omitempty"`
 	ServiceAccountName         string `json:"serviceAccountName,omitempty"`
 }
+
+// SetupWithManager creates and configures all components of the AppWrapper controller
+func SetupWithManager(ctx context.Context, mgr ctrl.Manager, config *AppWrapperConfig) error {
+	if err := WorkloadReconciler(
+		mgr.GetClient(),
+		mgr.GetEventRecorderFor("kueue"),
+		jobframework.WithManageJobsWithoutQueueName(config.ManageJobsWithoutQueueName),
+	).SetupWithManager(mgr); err != nil {
+		return fmt.Errorf("workload controller: %w", err)
+	}
+
+	if err := (&AppWrapperReconciler{
+		Client: mgr.GetClient(),
+		Scheme: mgr.GetScheme(),
+		Config: config,
+	}).SetupWithManager(mgr); err != nil {
+		return fmt.Errorf("appwrapper controller: %w", err)
+	}
+
+	if os.Getenv("ENABLE_WEBHOOKS") != "false" {
+		if err := (&AppWrapperWebhook{
+			Config: config,
+		}).SetupWebhookWithManager(mgr); err != nil {
+			return fmt.Errorf("appwrapper webhook: %w", err)
+		}
+	}
+
+	if err := jobframework.SetupWorkloadOwnerIndex(ctx, mgr.GetFieldIndexer(), GVK); err != nil {
+		return fmt.Errorf("appwrapper indexer: %w", err)
+	}
+
+	return nil
+}