Skip to content

Commit e0e4631

Browse files
SovietacedSovietaced
committed
Reuse more code
Signed-off-by: Jason Parraga <sovietaced@gmail.com>
1 parent deb5663 commit e0e4631

13 files changed

+314
-332
lines changed

‎internal/controller/common/helpers.go‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,17 @@ package common
33
import (
44
"context"
55

6+
"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
7+
68
"github.com/go-logr/logr"
79
k8serrors "k8s.io/apimachinery/pkg/api/errors"
810
"k8s.io/apimachinery/pkg/types"
911
"sigs.k8s.io/controller-runtime/pkg/client"
1012
)
1113

14+
// CleanupFunc is a function that will clean up additional resources which are not deleted by owner references.
15+
type CleanupFunc func(context.Context) error
16+
1217
// GetObject will get the object from Kubernetes and return if it is missing or an error.
1318
func GetObject(
1419
ctx context.Context,
@@ -27,3 +32,84 @@ func GetObject(
2732
}
2833
return false, nil
2934
}
35+
36+
// CheckAndHandleObjectDeletion handles the deletion of the resource by adding/removing the finalizer.
37+
// If the resource is being deleted, it will remove the finalizer.
38+
// If the resource is not being deleted, it will add the finalizer.
39+
// If finish is true, the reconciliation should finish early.
40+
func CheckAndHandleObjectDeletion(
41+
ctx context.Context,
42+
r client.Client,
43+
object client.Object,
44+
finalizer string,
45+
cleanupF CleanupFunc,
46+
logger logr.Logger,
47+
) (finish bool, err error) {
48+
logger = logger.WithValues("finalizer", finalizer)
49+
deletionTimestamp := object.GetDeletionTimestamp()
50+
if deletionTimestamp.IsZero() {
51+
// The object is not being deleted as deletionTimestamp.
52+
// In this case, we should add the finalizer if it is not already present.
53+
if err := addFinalizerIfNeeded(ctx, r, object, finalizer, logger); err != nil {
54+
return true, err
55+
}
56+
} else {
57+
// The object is being deleted so we should run the cleanup function if needed and remove the finalizer.
58+
return handleObjectDeletion(ctx, r, object, finalizer, cleanupF, logger)
59+
}
60+
// The object is not being deleted, continue reconciliation
61+
return false, nil
62+
}
63+
64+
// addFinalizerIfNeeded will add the finalizer to the object if it is not already present.
65+
func addFinalizerIfNeeded(
66+
ctx context.Context,
67+
client client.Client,
68+
object client.Object,
69+
finalizer string,
70+
logger logr.Logger,
71+
) error {
72+
if !controllerutil.ContainsFinalizer(object, finalizer) {
73+
logger.Info("Attaching cleanup finalizer because object does not have a deletion timestamp set")
74+
controllerutil.AddFinalizer(object, finalizer)
75+
return client.Update(ctx, object)
76+
}
77+
return nil
78+
}
79+
80+
func handleObjectDeletion(
81+
ctx context.Context,
82+
client client.Client,
83+
object client.Object,
84+
finalizer string,
85+
cleanupF CleanupFunc,
86+
logger logr.Logger,
87+
) (finish bool, err error) {
88+
deletionTimestamp := object.GetDeletionTimestamp()
89+
logger.Info(
90+
"Object is being deleted as it has a non-zero deletion timestamp set",
91+
"deletionTimestamp", deletionTimestamp,
92+
)
93+
logger.Info(
94+
"Namespace-scoped objects will be deleted by Kubernetes based on their OwnerReference",
95+
"deletionTimestamp", deletionTimestamp,
96+
)
97+
// The object is being deleted
98+
if controllerutil.ContainsFinalizer(object, finalizer) {
99+
// Run additional cleanup function if it is provided
100+
if cleanupF != nil {
101+
if err := cleanupF(ctx); err != nil {
102+
return true, err
103+
}
104+
}
105+
// Remove our finalizer from the list and update it.
106+
logger.Info("Removing cleanup finalizer from object")
107+
controllerutil.RemoveFinalizer(object, finalizer)
108+
if err := client.Update(ctx, object); err != nil {
109+
return true, err
110+
}
111+
}
112+
113+
// Stop reconciliation as the item is being deleted
114+
return true, nil
115+
}

‎internal/controller/common/helpers_test.go‎

Lines changed: 215 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,215 @@
1+
package common
2+
3+
import (
4+
"context"
5+
"testing"
6+
"time"
7+
8+
"github.com/go-logr/logr"
9+
"github.com/pkg/errors"
10+
"github.com/stretchr/testify/assert"
11+
k8serrors "k8s.io/apimachinery/pkg/api/errors"
12+
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
13+
"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
14+
"k8s.io/apimachinery/pkg/types"
15+
"k8s.io/apimachinery/pkg/util/uuid"
16+
"k8s.io/utils/ptr"
17+
"sigs.k8s.io/controller-runtime/pkg/client"
18+
"sigs.k8s.io/controller-runtime/pkg/client/fake"
19+
"sigs.k8s.io/controller-runtime/pkg/client/interceptor"
20+
"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
21+
22+
"k8s.io/apimachinery/pkg/runtime/schema"
23+
)
24+
25+
func TestCheckAndHandleResourceDeletion(t *testing.T) {
26+
t.Parallel()
27+
28+
ctx := context.Background()
29+
30+
// Create a logger (for tests, we use logr.Discard() to avoid actual logging)
31+
logger := logr.Discard()
32+
33+
// Table-driven test cases
34+
tests := []struct {
35+
name string
36+
deletionTimestamp *time.Time
37+
finalizerPresent bool
38+
expectFinalizer bool
39+
expectFinish bool
40+
expectError bool
41+
}{
42+
{
43+
name: "Object not being deleted, finalizer not present",
44+
deletionTimestamp: nil,
45+
finalizerPresent: false,
46+
expectFinalizer: true,
47+
expectFinish: false,
48+
expectError: false,
49+
},
50+
{
51+
name: "Object not being deleted, finalizer already present",
52+
deletionTimestamp: nil,
53+
finalizerPresent: true,
54+
expectFinalizer: true,
55+
expectFinish: false,
56+
expectError: false,
57+
},
58+
{
59+
name: "Object being deleted, finalizer present",
60+
deletionTimestamp: ptr.To(time.Now()),
61+
finalizerPresent: true,
62+
expectFinalizer: false,
63+
expectFinish: true,
64+
expectError: false,
65+
},
66+
}
67+
68+
for _, tc := range tests {
69+
t.Run(tc.name, func(t *testing.T) {
70+
// Create a fake object
71+
object := newTestObject(t)
72+
73+
// Set DeletionTimestamp if the test requires it
74+
if tc.deletionTimestamp != nil {
75+
object.SetDeletionTimestamp(&metav1.Time{Time: *tc.deletionTimestamp})
76+
}
77+
78+
// Add or remove finalizer based on the test case
79+
finalizer := "test.finalizer"
80+
if tc.finalizerPresent {
81+
controllerutil.AddFinalizer(object, finalizer)
82+
}
83+
84+
// Create a fake client
85+
fakeClient := fake.NewClientBuilder().WithObjects(object).Build()
86+
87+
cleanupFunc := func(ctx context.Context) error {
88+
if tc.deletionTimestamp == nil {
89+
t.Fatalf("cleanup function should not be called")
90+
}
91+
return nil
92+
}
93+
94+
// Call the function under test
95+
finish, err := CheckAndHandleObjectDeletion(ctx, fakeClient, object, finalizer, cleanupFunc, logger)
96+
97+
// Check for errors
98+
if (err != nil) != tc.expectError {
99+
t.Errorf("Expected error: %v, got: %v", tc.expectError, err)
100+
}
101+
102+
// Check if reconciliation should finish
103+
if finish != tc.expectFinish {
104+
t.Errorf("Expected finish: %v, got: %v", tc.expectFinish, finish)
105+
}
106+
107+
// Check if finalizer was added/removed as expected
108+
hasFinalizer := controllerutil.ContainsFinalizer(object, finalizer)
109+
if hasFinalizer != tc.expectFinalizer {
110+
t.Errorf("Expected finalizer: %v, got: %v", tc.expectFinalizer, hasFinalizer)
111+
}
112+
})
113+
}
114+
}
115+
116+
func TestGetObjectFromCache(t *testing.T) {
117+
t.Parallel()
118+
119+
ctx := context.Background()
120+
121+
// No-op logger for testing
122+
logger := logr.Discard()
123+
124+
// Test cases
125+
tests := []struct {
126+
name string
127+
objectExists bool
128+
expectMiss bool
129+
expectError bool
130+
returnError error
131+
}{
132+
{
133+
name: "Object exists in cache",
134+
objectExists: true,
135+
returnError: nil,
136+
expectMiss: false,
137+
expectError: false,
138+
},
139+
{
140+
name: "Object not found in cache",
141+
objectExists: false,
142+
returnError: k8serrors.NewNotFound(newTestGroupResource(t), "test-resource"),
143+
expectMiss: true,
144+
expectError: false,
145+
},
146+
{
147+
name: "Error while fetching from cache",
148+
objectExists: false,
149+
returnError: errors.New("some network error"),
150+
expectMiss: true,
151+
expectError: true,
152+
},
153+
}
154+
155+
for _, tc := range tests {
156+
t.Run(tc.name, func(t *testing.T) {
157+
// Create a fake client with the expected error behavior
158+
clientBuilder := fake.NewClientBuilder()
159+
160+
// Create a fake object
161+
object := newTestObject(t)
162+
163+
if tc.objectExists {
164+
clientBuilder.WithObjects(object)
165+
}
166+
167+
if tc.expectError {
168+
clientBuilder.WithInterceptorFuncs(interceptor.Funcs{
169+
Get: func(ctx context.Context, client client.WithWatch, key client.ObjectKey, obj client.Object, opts ...client.GetOption) error {
170+
return tc.returnError
171+
},
172+
})
173+
}
174+
175+
fakeClient := clientBuilder.Build()
176+
177+
// Call the function under test
178+
namespacedName := types.NamespacedName{Name: "test-resource", Namespace: "default"}
179+
miss, err := GetObject(ctx, fakeClient, object, namespacedName, logger)
180+
if tc.expectError {
181+
assert.ErrorIs(t, err, tc.returnError)
182+
} else {
183+
assert.NoError(t, err)
184+
}
185+
assert.Equal(t, miss, tc.expectMiss)
186+
})
187+
}
188+
}
189+
190+
// newTestObject returns a test unstructured.Unstructured object only to be used in tests.
191+
func newTestObject(t *testing.T) *unstructured.Unstructured {
192+
object := &unstructured.Unstructured{}
193+
object.SetNamespace("default")
194+
object.SetName("test-resource")
195+
object.SetUID(uuid.NewUUID())
196+
object.SetGroupVersionKind(newTestGroupVersionKind(t))
197+
return object
198+
}
199+
200+
// newTestGroupVersionKind returns a test schema.GroupVersionKind only to be used in tests.
201+
func newTestGroupVersionKind(t *testing.T) schema.GroupVersionKind {
202+
return schema.GroupVersionKind{
203+
Group: "test.group",
204+
Version: "v1",
205+
Kind: "TestKind",
206+
}
207+
}
208+
209+
// newTestGroupResource returns a test schema.GroupResource only to be used in tests.
210+
func newTestGroupResource(t *testing.T) schema.GroupResource {
211+
return schema.GroupResource{
212+
Group: "test.group",
213+
Resource: "test-resource",
214+
}
215+
}

‎internal/controller/core/queue_controller.go‎

Lines changed: 5 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,6 @@ import (
2828
"google.golang.org/grpc/codes"
2929
"google.golang.org/grpc/status"
3030
"google.golang.org/protobuf/testing/protocmp"
31-
"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
32-
3331
"k8s.io/apimachinery/pkg/runtime"
3432
ctrl "sigs.k8s.io/controller-runtime"
3533
"sigs.k8s.io/controller-runtime/pkg/client"
@@ -63,42 +61,13 @@ func (r *QueueReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl
6361
return ctrl.Result{}, err
6462
}
6563

66-
// examine DeletionTimestamp to determine if object is under deletion
67-
deletionTimestamp := queue.ObjectMeta.DeletionTimestamp
68-
// examine DeletionTimestamp to determine if object is under deletion
69-
if deletionTimestamp.IsZero() {
70-
// The object is not being deleted, so if it does not have our finalizer,
71-
// then lets add the finalizer and update the object. This is equivalent
72-
// registering our finalizer.
73-
if !controllerutil.ContainsFinalizer(&queue, operatorFinalizer) {
74-
logger.Info("Attaching finalizer to Lookout object", "finalizer", operatorFinalizer)
75-
controllerutil.AddFinalizer(&queue, operatorFinalizer)
76-
if err := r.Update(ctx, &queue); err != nil {
77-
return ctrl.Result{}, err
78-
}
79-
}
80-
} else {
81-
logger.Info("Queue object is being deleted", "finalizer", operatorFinalizer)
82-
logger.Info("Namespace-scoped resources will be deleted by Kubernetes based on their OwnerReference")
83-
// The object is being deleted
84-
if controllerutil.ContainsFinalizer(&queue, operatorFinalizer) {
85-
// remove our finalizer from the list and update it.
86-
logger.Info("Removing finalizer from Queue object", "finalizer", operatorFinalizer)
87-
controllerutil.RemoveFinalizer(&queue, operatorFinalizer)
88-
if err := r.Update(ctx, &queue); err != nil {
89-
return ctrl.Result{}, err
90-
}
91-
}
92-
93-
if err := r.deleteQueue(ctx, queue); err != nil {
94-
return ctrl.Result{}, err
95-
}
96-
97-
// Stop reconciliation as the item is being deleted
98-
return ctrl.Result{}, nil
64+
cleanupF := func(ctx context.Context) error { return r.deleteQueue(ctx, queue) }
65+
finish, err := common.CheckAndHandleObjectDeletion(ctx, r.Client, &queue, operatorFinalizer, cleanupF, logger)
66+
if err != nil || finish {
67+
return ctrl.Result{}, err
9968
}
10069

101-
err := r.upsertQueue(ctx, queue)
70+
err = r.upsertQueue(ctx, queue)
10271
if err != nil {
10372
return ctrl.Result{}, err
10473
}

‎internal/controller/install/armadaserver_controller.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,7 @@ func (r *ArmadaServerReconciler) Reconcile(ctx context.Context, req ctrl.Request
8989
cleanupF := func(ctx context.Context) error {
9090
return r.deleteExternalResources(ctx, components, logger)
9191
}
92-
finish, err := checkAndHandleObjectDeletion(ctx, r.Client, &server, operatorFinalizer, cleanupF, logger)
92+
finish, err := common.CheckAndHandleObjectDeletion(ctx, r.Client, &server, operatorFinalizer, cleanupF, logger)
9393
if err != nil || finish {
9494
return ctrl.Result{}, err
9595
}

‎internal/controller/install/binoculars_controller.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ func (r *BinocularsReconciler) Reconcile(ctx context.Context, req ctrl.Request)
8585
cleanupF := func(ctx context.Context) error {
8686
return r.deleteExternalResources(ctx, components)
8787
}
88-
finish, err := checkAndHandleObjectDeletion(ctx, r.Client, &binoculars, operatorFinalizer, cleanupF, logger)
88+
finish, err := common.CheckAndHandleObjectDeletion(ctx, r.Client, &binoculars, operatorFinalizer, cleanupF, logger)
8989
if err != nil || finish {
9090
return ctrl.Result{}, err
9191
}

0 commit comments

Comments
 (0)