Skip to content

Commit a8a8a4d

Browse files
ASverdlovASverdlov
committed
Test WorkloadPriorityClassReconciler
1 parent 40cbe0d commit a8a8a4d

File tree

1 file changed

+295
-0
lines changed

1 file changed

+295
-0
lines changed

‎pkg/controller/core/workloadpriorityclass_controller_test.go‎

Lines changed: 295 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,295 @@
1+
/*
2+
Copyright The Kubernetes Authors.
3+
4+
Licensed under the Apache License, Version 2.0 (the "License");
5+
you may not use this file except in compliance with the License.
6+
You may obtain a copy of the License at
7+
8+
http://www.apache.org/licenses/LICENSE-2.0
9+
10+
Unless required by applicable law or agreed to in writing, software
11+
distributed under the License is distributed on an "AS IS" BASIS,
12+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
See the License for the specific language governing permissions and
14+
limitations under the License.
15+
*/
16+
17+
package core
18+
19+
import (
20+
"context"
21+
stderrors "errors"
22+
"testing"
23+
24+
"github.com/google/go-cmp/cmp"
25+
"k8s.io/apimachinery/pkg/api/errors"
26+
"k8s.io/apimachinery/pkg/types"
27+
"sigs.k8s.io/controller-runtime/pkg/client"
28+
"sigs.k8s.io/controller-runtime/pkg/client/interceptor"
29+
"sigs.k8s.io/controller-runtime/pkg/event"
30+
"sigs.k8s.io/controller-runtime/pkg/reconcile"
31+
32+
kueue "sigs.k8s.io/kueue/apis/kueue/v1beta2"
33+
"sigs.k8s.io/kueue/pkg/controller/core/indexer"
34+
utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
35+
utiltestingapi "sigs.k8s.io/kueue/pkg/util/testing/v1beta2"
36+
)
37+
38+
func TestWorkloadPriorityClassPredicates(t *testing.T) {
39+
cases := map[string]struct {
40+
eventType string
41+
oldWPC *kueue.WorkloadPriorityClass
42+
newWPC *kueue.WorkloadPriorityClass
43+
want bool
44+
}{
45+
"create event should trigger reconcile": {
46+
eventType: "create",
47+
newWPC: utiltestingapi.MakeWorkloadPriorityClass("test").PriorityValue(100).Obj(),
48+
want: true,
49+
},
50+
"delete event should not trigger reconcile": {
51+
eventType: "delete",
52+
oldWPC: utiltestingapi.MakeWorkloadPriorityClass("test").PriorityValue(100).Obj(),
53+
want: false,
54+
},
55+
"update event with changed priority should trigger reconcile": {
56+
eventType: "update",
57+
oldWPC: utiltestingapi.MakeWorkloadPriorityClass("test").PriorityValue(100).Obj(),
58+
newWPC: utiltestingapi.MakeWorkloadPriorityClass("test").PriorityValue(200).Obj(),
59+
want: true,
60+
},
61+
"update event with unchanged priority should not trigger reconcile": {
62+
eventType: "update",
63+
oldWPC: utiltestingapi.MakeWorkloadPriorityClass("test").PriorityValue(100).Obj(),
64+
newWPC: utiltestingapi.MakeWorkloadPriorityClass("test").PriorityValue(100).Obj(),
65+
want: false,
66+
},
67+
"generic event should not trigger reconcile": {
68+
eventType: "generic",
69+
newWPC: utiltestingapi.MakeWorkloadPriorityClass("test").PriorityValue(100).Obj(),
70+
want: false,
71+
},
72+
}
73+
74+
for name, tc := range cases {
75+
t.Run(name, func(t *testing.T) {
76+
reconciler := NewWorkloadPriorityClassReconciler(nil, nil)
77+
var got bool
78+
79+
switch tc.eventType {
80+
case "create":
81+
got = reconciler.Create(event.TypedCreateEvent[*kueue.WorkloadPriorityClass]{Object: tc.newWPC})
82+
case "delete":
83+
got = reconciler.Delete(event.TypedDeleteEvent[*kueue.WorkloadPriorityClass]{Object: tc.oldWPC})
84+
case "update":
85+
got = reconciler.Update(event.TypedUpdateEvent[*kueue.WorkloadPriorityClass]{
86+
ObjectOld: tc.oldWPC,
87+
ObjectNew: tc.newWPC,
88+
})
89+
case "generic":
90+
got = reconciler.Generic(event.TypedGenericEvent[*kueue.WorkloadPriorityClass]{Object: tc.newWPC})
91+
}
92+
93+
if got != tc.want {
94+
t.Errorf("got %v, want %v", got, tc.want)
95+
}
96+
})
97+
}
98+
}
99+
100+
func TestWorkloadPriorityClassReconcile(t *testing.T) {
101+
cases := map[string]struct {
102+
wpc *kueue.WorkloadPriorityClass
103+
workloads []kueue.Workload
104+
wantWorkloads []kueue.Workload
105+
wantError bool
106+
clientFuncs *interceptor.Funcs
107+
}{
108+
"reconcile updates workload priority when WPC priority changes": {
109+
wpc: utiltestingapi.MakeWorkloadPriorityClass("high").PriorityValue(1000).Obj(),
110+
workloads: []kueue.Workload{
111+
*utiltestingapi.MakeWorkload("wl1", "default").
112+
Priority(100).
113+
WorkloadPriorityClassRef("high").
114+
Obj(),
115+
},
116+
wantWorkloads: []kueue.Workload{
117+
*utiltestingapi.MakeWorkload("wl1", "default").
118+
Priority(1000).
119+
WorkloadPriorityClassRef("high").
120+
Obj(),
121+
},
122+
},
123+
"reconcile updates multiple workloads": {
124+
wpc: utiltestingapi.MakeWorkloadPriorityClass("high").PriorityValue(1000).Obj(),
125+
workloads: []kueue.Workload{
126+
*utiltestingapi.MakeWorkload("wl1", "default").
127+
Priority(100).
128+
WorkloadPriorityClassRef("high").
129+
Obj(),
130+
*utiltestingapi.MakeWorkload("wl2", "default").
131+
Priority(200).
132+
WorkloadPriorityClassRef("high").
133+
Obj(),
134+
},
135+
wantWorkloads: []kueue.Workload{
136+
*utiltestingapi.MakeWorkload("wl1", "default").
137+
Priority(1000).
138+
WorkloadPriorityClassRef("high").
139+
Obj(),
140+
*utiltestingapi.MakeWorkload("wl2", "default").
141+
Priority(1000).
142+
WorkloadPriorityClassRef("high").
143+
Obj(),
144+
},
145+
},
146+
"reconcile skips workloads with up-to-date priority": {
147+
wpc: utiltestingapi.MakeWorkloadPriorityClass("high").PriorityValue(1000).Obj(),
148+
workloads: []kueue.Workload{
149+
*utiltestingapi.MakeWorkload("wl1", "default").
150+
Priority(1000).
151+
WorkloadPriorityClassRef("high").
152+
Obj(),
153+
},
154+
wantWorkloads: []kueue.Workload{
155+
*utiltestingapi.MakeWorkload("wl1", "default").
156+
Priority(1000).
157+
WorkloadPriorityClassRef("high").
158+
Obj(),
159+
},
160+
},
161+
"reconcile succeeds when no workloads use the WPC": {
162+
wpc: utiltestingapi.MakeWorkloadPriorityClass("high").PriorityValue(1000).Obj(),
163+
workloads: []kueue.Workload{},
164+
wantWorkloads: []kueue.Workload{},
165+
},
166+
"reconcile handles workload not found error": {
167+
wpc: utiltestingapi.MakeWorkloadPriorityClass("high").PriorityValue(1000).Obj(),
168+
workloads: []kueue.Workload{
169+
*utiltestingapi.MakeWorkload("wl1", "default").
170+
Priority(100).
171+
WorkloadPriorityClassRef("high").
172+
Obj(),
173+
},
174+
clientFuncs: &interceptor.Funcs{
175+
Update: func(ctx context.Context, client client.WithWatch, obj client.Object, opts ...client.UpdateOption) error {
176+
return errors.NewNotFound(kueue.Resource("workload"), "wl1")
177+
},
178+
},
179+
wantWorkloads: []kueue.Workload{
180+
*utiltestingapi.MakeWorkload("wl1", "default").
181+
Priority(100).
182+
WorkloadPriorityClassRef("high").
183+
Obj(),
184+
},
185+
},
186+
"reconcile returns error when update fails": {
187+
wpc: utiltestingapi.MakeWorkloadPriorityClass("high").PriorityValue(1000).Obj(),
188+
workloads: []kueue.Workload{
189+
*utiltestingapi.MakeWorkload("wl1", "default").
190+
Priority(100).
191+
WorkloadPriorityClassRef("high").
192+
Obj(),
193+
},
194+
clientFuncs: &interceptor.Funcs{
195+
Update: func(ctx context.Context, client client.WithWatch, obj client.Object, opts ...client.UpdateOption) error {
196+
return stderrors.New("update failed")
197+
},
198+
},
199+
wantError: true,
200+
wantWorkloads: []kueue.Workload{
201+
*utiltestingapi.MakeWorkload("wl1", "default").
202+
Priority(100).
203+
WorkloadPriorityClassRef("high").
204+
Obj(),
205+
},
206+
},
207+
"reconcile handles partial update failures": {
208+
wpc: utiltestingapi.MakeWorkloadPriorityClass("high").PriorityValue(1000).Obj(),
209+
workloads: []kueue.Workload{
210+
*utiltestingapi.MakeWorkload("wl1", "default").
211+
Priority(100).
212+
WorkloadPriorityClassRef("high").
213+
Obj(),
214+
*utiltestingapi.MakeWorkload("wl2", "default").
215+
Priority(200).
216+
WorkloadPriorityClassRef("high").
217+
Obj(),
218+
},
219+
clientFuncs: &interceptor.Funcs{
220+
Update: func(ctx context.Context, client client.WithWatch, obj client.Object, opts ...client.UpdateOption) error {
221+
wl := obj.(*kueue.Workload)
222+
if wl.Name == "wl2" {
223+
return stderrors.New("update failed for wl2")
224+
}
225+
return client.Update(ctx, obj, opts...)
226+
},
227+
},
228+
wantError: true,
229+
wantWorkloads: []kueue.Workload{
230+
*utiltestingapi.MakeWorkload("wl1", "default").
231+
Priority(1000).
232+
WorkloadPriorityClassRef("high").
233+
Obj(),
234+
*utiltestingapi.MakeWorkload("wl2", "default").
235+
Priority(200).
236+
WorkloadPriorityClassRef("high").
237+
Obj(),
238+
},
239+
},
240+
"reconcile handles WPC not found": {
241+
wpc: utiltestingapi.MakeWorkloadPriorityClass("high").PriorityValue(1000).Obj(),
242+
workloads: []kueue.Workload{},
243+
wantWorkloads: []kueue.Workload{},
244+
clientFuncs: &interceptor.Funcs{
245+
Get: func(ctx context.Context, client client.WithWatch, key client.ObjectKey, obj client.Object, opts ...client.GetOption) error {
246+
return errors.NewNotFound(kueue.Resource("workloadpriorityclass"), key.Name)
247+
},
248+
},
249+
},
250+
}
251+
252+
for name, tc := range cases {
253+
t.Run(name, func(t *testing.T) {
254+
ctx := context.Background()
255+
256+
builder := utiltesting.NewClientBuilder().
257+
WithObjects(tc.wpc).
258+
WithIndex(&kueue.Workload{}, indexer.WorkloadPriorityClassKey, indexer.IndexWorkloadPriorityClass).
259+
WithStatusSubresource(&kueue.Workload{})
260+
for i := range tc.workloads {
261+
builder = builder.WithObjects(&tc.workloads[i])
262+
}
263+
if tc.clientFuncs != nil {
264+
builder = builder.WithInterceptorFuncs(*tc.clientFuncs)
265+
}
266+
k8sClient := builder.Build()
267+
268+
reconciler := NewWorkloadPriorityClassReconciler(k8sClient, nil)
269+
req := reconcile.Request{
270+
NamespacedName: types.NamespacedName{
271+
Name: tc.wpc.Name,
272+
},
273+
}
274+
275+
_, gotErr := reconciler.Reconcile(ctx, req)
276+
277+
if tc.wantError && gotErr == nil {
278+
t.Errorf("expected error but got nil")
279+
} else if !tc.wantError && gotErr != nil {
280+
t.Errorf("unexpected error: %v", gotErr)
281+
}
282+
// Verify workloads are in the expected state
283+
for _, wantWl := range tc.wantWorkloads {
284+
gotWl := &kueue.Workload{}
285+
err := k8sClient.Get(ctx, types.NamespacedName{Name: wantWl.Name, Namespace: wantWl.Namespace}, gotWl)
286+
if err != nil {
287+
t.Fatalf("failed to get workload %s: %v", wantWl.Name, err)
288+
}
289+
if diff := cmp.Diff(wantWl.Spec.Priority, gotWl.Spec.Priority); diff != "" {
290+
t.Errorf("workload %s priority mismatch (-want +got):\n%s", wantWl.Name, diff)
291+
}
292+
}
293+
})
294+
}
295+
}

0 commit comments

Comments
 (0)