Add unit tests for pkg/shard/controller (#486)

* Prefactor: extract `matchers.BeFunc`

* Prefactor: simplify instantiating `shardcontroller.Reconciler`

* Add unit tests for `pkg/shard/controller`

Author: timebertt

pkg/shard/controller/builder.go

@@ -23,8 +23,6 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/manager"
 	"sigs.k8s.io/controller-runtime/pkg/reconcile"
-
-	shardingv1alpha1 "github.com/timebertt/kubernetes-controller-sharding/pkg/apis/sharding/v1alpha1"
 )
 
 // Builder can build a sharded reconciler.
@@ -42,9 +40,7 @@ type Builder struct {
 // reconciler that takes care of the sharding-related logic and calls the delegate reconciler whenever the shard is
 // responsible for reconciling an object.
 func NewShardedReconciler(mgr manager.Manager) *Builder {
-	return &Builder{
-		client: mgr.GetClient(),
-	}
+	return (&Builder{}).WithClient(mgr.GetClient())
 }
 
 // For sets the object kind being reconciled by the reconciler.
@@ -105,11 +101,10 @@ func (b *Builder) Build(r reconcile.Reconciler) (reconcile.Reconciler, error) {
 	}
 
 	return &Reconciler{
-		Object:     b.object,
-		Client:     b.client,
-		ShardName:  b.shardName,
-		LabelShard: shardingv1alpha1.LabelShard(b.controllerRingName),
-		LabelDrain: shardingv1alpha1.LabelDrain(b.controllerRingName),
-		Do:         r,
+		Object:             b.object,
+		Client:             b.client,
+		ControllerRingName: b.controllerRingName,
+		ShardName:          b.shardName,
+		Do:                 r,
 	}, nil
 }

pkg/shard/controller/builder_test.go

@@ -0,0 +1,154 @@
+/*
+Copyright 2025 Tim Ebert.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package controller_test
+
+import (
+	"context"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+	corev1 "k8s.io/api/core/v1"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	fakeclient "sigs.k8s.io/controller-runtime/pkg/client/fake"
+	"sigs.k8s.io/controller-runtime/pkg/manager"
+	"sigs.k8s.io/controller-runtime/pkg/reconcile"
+
+	. "github.com/timebertt/kubernetes-controller-sharding/pkg/shard/controller"
+	. "github.com/timebertt/kubernetes-controller-sharding/pkg/utils/test/matchers"
+)
+
+var _ = Describe("#Builder", func() {
+	var (
+		mgr fakeManager
+		b   *Builder
+
+		obj                *corev1.Pod
+		controllerRingName string
+		shardName          string
+		client             client.Client
+		r                  reconcile.Reconciler
+	)
+
+	BeforeEach(func() {
+		obj = &corev1.Pod{}
+		controllerRingName = "operator"
+		shardName = "operator-0"
+		client = fakeclient.NewFakeClient()
+		r = reconcile.Func(func(context.Context, reconcile.Request) (reconcile.Result, error) {
+			return reconcile.Result{}, nil
+		})
+	})
+
+	JustBeforeEach(func() {
+		mgr = fakeManager{Client: client}
+		b = NewShardedReconciler(mgr).
+			For(obj).
+			InControllerRing(controllerRingName).
+			WithShardName(shardName)
+	})
+
+	Describe("#For", func() {
+		It("should complain about calling For twice", func() {
+			Expect(b.For(obj).Build(r)).Error().To(MatchError("must not call For() more than once"))
+		})
+
+		It("should complain about not calling For", func() {
+			b = NewShardedReconciler(mgr).
+				InControllerRing(controllerRingName).
+				WithShardName(shardName)
+			Expect(b.Build(r)).Error().To(MatchError("missing object kind, must call to For()"))
+		})
+	})
+
+	Describe("#WithClient", func() {
+		It("should use the custom client instead of the manager's client", func() {
+			customClient := fakeclient.NewFakeClient()
+			Expect(b.WithClient(customClient).Build(r)).To(
+				HaveField("Client", BeIdenticalTo(customClient)),
+			)
+		})
+
+		It("should complain about missing client", func() {
+			Expect(b.WithClient(nil).Build(r)).Error().To(MatchError("missing client"))
+		})
+	})
+
+	Describe("#InControllerRing", func() {
+		It("should complain about missing ControllerRing name", func() {
+			Expect(b.InControllerRing("").Build(r)).Error().To(MatchError("missing ControllerRing name"))
+		})
+	})
+
+	Describe("#WithShardName", func() {
+		It("should complain about missing shard name", func() {
+			Expect(b.WithShardName("").Build(r)).Error().To(MatchError("missing shard name"))
+		})
+	})
+
+	Describe("#Build", func() {
+		It("should complain about nil reconciler", func() {
+			Expect(b.Build(nil)).Error().To(MatchError("must provide a non-nil Reconciler"))
+		})
+
+		It("should correctly set up the Reconciler", func() {
+			shardReconciler, err := b.Build(r)
+			Expect(err).NotTo(HaveOccurred())
+
+			reconciler, ok := shardReconciler.(*Reconciler)
+			Expect(ok).To(BeTrue())
+
+			Expect(reconciler.Object).To(BeAssignableToTypeOf(obj))
+			Expect(reconciler.Client).To(BeIdenticalTo(client))
+			Expect(reconciler.ControllerRingName).To(Equal(controllerRingName))
+			Expect(reconciler.ShardName).To(Equal(shardName))
+			Expect(reconciler.Do).To(BeFunc(r))
+		})
+	})
+
+	Describe("#MustBuild", func() {
+		It("should panic for nil reconciler", func() {
+			Expect(func() {
+				b.MustBuild(nil)
+			}).To(Panic())
+		})
+
+		It("should correctly set up the Reconciler", func() {
+			var shardReconciler reconcile.Reconciler
+			Expect(func() {
+				shardReconciler = b.MustBuild(r)
+			}).NotTo(Panic())
+
+			reconciler, ok := shardReconciler.(*Reconciler)
+			Expect(ok).To(BeTrue())
+
+			Expect(reconciler.Object).To(BeAssignableToTypeOf(obj))
+			Expect(reconciler.Client).To(BeIdenticalTo(client))
+			Expect(reconciler.ControllerRingName).To(Equal(controllerRingName))
+			Expect(reconciler.ShardName).To(Equal(shardName))
+			Expect(reconciler.Do).To(BeFunc(r))
+		})
+	})
+})
+
+type fakeManager struct {
+	manager.Manager
+	Client client.Client
+}
+
+func (f fakeManager) GetClient() client.Client {
+	return f.Client
+}

pkg/shard/controller/controller_suite_test.go

@@ -0,0 +1,29 @@
+/*
+Copyright 2025 Tim Ebert.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package controller_test
+
+import (
+	"testing"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+func TestController(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "Shard Controller Suite")
+}

pkg/shard/controller/predicate_test.go

@@ -0,0 +1,129 @@
+/*
+Copyright 2025 Tim Ebert.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package controller_test
+
+import (
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/event"
+	"sigs.k8s.io/controller-runtime/pkg/predicate"
+
+	. "github.com/timebertt/kubernetes-controller-sharding/pkg/shard/controller"
+)
+
+var _ = Describe("#Predicate", func() {
+	var (
+		controllerRingName string
+		shardName          string
+
+		mainPredicate, p predicate.Predicate
+
+		obj, objOld *corev1.Pod
+	)
+
+	BeforeEach(func() {
+		controllerRingName = "operator"
+		shardName = "operator-0"
+
+		obj = &corev1.Pod{}
+	})
+
+	JustBeforeEach(func() {
+		p = Predicate(controllerRingName, shardName, mainPredicate)
+	})
+
+	When("main predicate returns false", func() {
+		BeforeEach(func() {
+			mainPredicate = predicate.NewPredicateFuncs(func(client.Object) bool {
+				return false
+			})
+		})
+
+		It("should handle drained objects", func() {
+			metav1.SetMetaDataLabel(&obj.ObjectMeta, "drain.alpha.sharding.timebertt.dev/"+controllerRingName, "true")
+			objOld = obj.DeepCopy()
+
+			Expect(p.Create(event.CreateEvent{Object: obj})).To(BeTrue())
+			Expect(p.Update(event.UpdateEvent{ObjectOld: objOld, ObjectNew: obj})).To(BeTrue())
+			Expect(p.Delete(event.DeleteEvent{Object: obj})).To(BeTrue())
+			Expect(p.Generic(event.GenericEvent{Object: obj})).To(BeTrue())
+		})
+
+		It("should handle assigned and drained objects", func() {
+			metav1.SetMetaDataLabel(&obj.ObjectMeta, "shard.alpha.sharding.timebertt.dev/"+controllerRingName, shardName)
+			metav1.SetMetaDataLabel(&obj.ObjectMeta, "drain.alpha.sharding.timebertt.dev/"+controllerRingName, "true")
+			objOld = obj.DeepCopy()
+
+			Expect(p.Create(event.CreateEvent{Object: obj})).To(BeTrue())
+			Expect(p.Update(event.UpdateEvent{ObjectOld: objOld, ObjectNew: obj})).To(BeTrue())
+			Expect(p.Delete(event.DeleteEvent{Object: obj})).To(BeTrue())
+			Expect(p.Generic(event.GenericEvent{Object: obj})).To(BeTrue())
+		})
+
+		It("should not handle assigned objects", func() {
+			metav1.SetMetaDataLabel(&obj.ObjectMeta, "shard.alpha.sharding.timebertt.dev/"+controllerRingName, shardName)
+			objOld = obj.DeepCopy()
+
+			Expect(p.Create(event.CreateEvent{Object: obj})).To(BeFalse())
+			Expect(p.Update(event.UpdateEvent{ObjectOld: objOld, ObjectNew: obj})).To(BeFalse())
+			Expect(p.Delete(event.DeleteEvent{Object: obj})).To(BeFalse())
+			Expect(p.Generic(event.GenericEvent{Object: obj})).To(BeFalse())
+		})
+	})
+
+	When("main predicate returns true", func() {
+		BeforeEach(func() {
+			mainPredicate = predicate.NewPredicateFuncs(func(client.Object) bool {
+				return true
+			})
+		})
+
+		It("should handle drained objects", func() {
+			metav1.SetMetaDataLabel(&obj.ObjectMeta, "drain.alpha.sharding.timebertt.dev/"+controllerRingName, "true")
+			objOld = obj.DeepCopy()
+
+			Expect(p.Create(event.CreateEvent{Object: obj})).To(BeTrue())
+			Expect(p.Update(event.UpdateEvent{ObjectOld: objOld, ObjectNew: obj})).To(BeTrue())
+			Expect(p.Delete(event.DeleteEvent{Object: obj})).To(BeTrue())
+			Expect(p.Generic(event.GenericEvent{Object: obj})).To(BeTrue())
+		})
+
+		It("should handle assigned objects", func() {
+			metav1.SetMetaDataLabel(&obj.ObjectMeta, "shard.alpha.sharding.timebertt.dev/"+controllerRingName, shardName)
+			objOld = obj.DeepCopy()
+
+			Expect(p.Create(event.CreateEvent{Object: obj})).To(BeTrue())
+			Expect(p.Update(event.UpdateEvent{ObjectOld: objOld, ObjectNew: obj})).To(BeTrue())
+			Expect(p.Delete(event.DeleteEvent{Object: obj})).To(BeTrue())
+			Expect(p.Generic(event.GenericEvent{Object: obj})).To(BeTrue())
+		})
+
+		It("should handle assigned and drained objects", func() {
+			metav1.SetMetaDataLabel(&obj.ObjectMeta, "shard.alpha.sharding.timebertt.dev/"+controllerRingName, shardName)
+			metav1.SetMetaDataLabel(&obj.ObjectMeta, "drain.alpha.sharding.timebertt.dev/"+controllerRingName, "true")
+			objOld = obj.DeepCopy()
+
+			Expect(p.Create(event.CreateEvent{Object: obj})).To(BeTrue())
+			Expect(p.Update(event.UpdateEvent{ObjectOld: objOld, ObjectNew: obj})).To(BeTrue())
+			Expect(p.Delete(event.DeleteEvent{Object: obj})).To(BeTrue())
+			Expect(p.Generic(event.GenericEvent{Object: obj})).To(BeTrue())
+		})
+	})
+})