Skip to content

Commit c1a7dd3

Browse files
harcheharche
authored
Add DRA ResourceClaim wrappers for test infrastructure (#7300)
This commit introduces test helper wrappers for DRA ResourceClaim and ResourceClaimTemplate objects to improve test readability and reduce boilerplate. Signed-off-by: Harshal Patil <[email protected]>
1 parent bb2b22f commit c1a7dd3

File tree

5 files changed

+349
-287
lines changed

5 files changed

+349
-287
lines changed

‎pkg/controller/core/workload_controller_test.go‎

Lines changed: 20 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -427,20 +427,11 @@ func TestReconcile(t *testing.T) {
427427
ResourceClaim("gpu", "rc1").
428428
Obj()).
429429
Obj(),
430-
resourceClaims: []*resourcev1.ResourceClaim{{
431-
ObjectMeta: metav1.ObjectMeta{Name: "rc1", Namespace: "ns"},
432-
Spec: resourcev1.ResourceClaimSpec{
433-
Devices: resourcev1.DeviceClaim{
434-
Requests: []resourcev1.DeviceRequest{{
435-
Exactly: &resourcev1.ExactDeviceRequest{
436-
DeviceClassName: "gpu.example.com",
437-
AllocationMode: resourcev1.DeviceAllocationModeExactCount,
438-
Count: 1,
439-
},
440-
}},
441-
},
442-
},
443-
}},
430+
resourceClaims: []*resourcev1.ResourceClaim{
431+
utiltesting.MakeResourceClaim("rc1", "ns").
432+
DeviceRequest("", "gpu.example.com", 1).
433+
Obj(),
434+
},
444435
cq: utiltestingapi.MakeClusterQueue("cq").
445436
ResourceGroup(
446437
*utiltestingapi.MakeFlavorQuotas("flavor1").
@@ -477,23 +468,11 @@ func TestReconcile(t *testing.T) {
477468
ResourceClaimTemplate("gpu", "gpu-template").
478469
Obj()).
479470
Obj(),
480-
resourceClaimTemplates: []*resourcev1.ResourceClaimTemplate{{
481-
ObjectMeta: metav1.ObjectMeta{Name: "gpu-template", Namespace: "ns"},
482-
Spec: resourcev1.ResourceClaimTemplateSpec{
483-
Spec: resourcev1.ResourceClaimSpec{
484-
Devices: resourcev1.DeviceClaim{
485-
Requests: []resourcev1.DeviceRequest{{
486-
Name: "gpu-request",
487-
Exactly: &resourcev1.ExactDeviceRequest{
488-
DeviceClassName: "gpu.example.com",
489-
AllocationMode: resourcev1.DeviceAllocationModeExactCount,
490-
Count: 1,
491-
},
492-
}},
493-
},
494-
},
495-
},
496-
}},
471+
resourceClaimTemplates: []*resourcev1.ResourceClaimTemplate{
472+
utiltesting.MakeResourceClaimTemplate("gpu-template", "ns").
473+
DeviceRequest("gpu-request", "gpu.example.com", 1).
474+
Obj(),
475+
},
497476
cq: utiltestingapi.MakeClusterQueue("cq").
498477
ResourceGroup(
499478
*utiltestingapi.MakeFlavorQuotas("flavor1").
@@ -524,23 +503,11 @@ func TestReconcile(t *testing.T) {
524503
ResourceClaimTemplate("gpu", "gpu-template").
525504
Obj()).
526505
Obj(),
527-
resourceClaimTemplates: []*resourcev1.ResourceClaimTemplate{{
528-
ObjectMeta: metav1.ObjectMeta{Name: "gpu-template", Namespace: "ns"},
529-
Spec: resourcev1.ResourceClaimTemplateSpec{
530-
Spec: resourcev1.ResourceClaimSpec{
531-
Devices: resourcev1.DeviceClaim{
532-
Requests: []resourcev1.DeviceRequest{{
533-
Name: "gpu-request",
534-
Exactly: &resourcev1.ExactDeviceRequest{
535-
DeviceClassName: "gpu.example.com",
536-
AllocationMode: resourcev1.DeviceAllocationModeExactCount,
537-
Count: 2,
538-
},
539-
}},
540-
},
541-
},
542-
},
543-
}},
506+
resourceClaimTemplates: []*resourcev1.ResourceClaimTemplate{
507+
utiltesting.MakeResourceClaimTemplate("gpu-template", "ns").
508+
DeviceRequest("gpu-request", "gpu.example.com", 2).
509+
Obj(),
510+
},
544511
cq: utiltestingapi.MakeClusterQueue("cq").
545512
ResourceGroup(
546513
*utiltestingapi.MakeFlavorQuotas("flavor1").
@@ -569,23 +536,11 @@ func TestReconcile(t *testing.T) {
569536
ResourceClaimTemplate("gpu", "gpu-template").
570537
Obj()).
571538
Obj(),
572-
resourceClaimTemplates: []*resourcev1.ResourceClaimTemplate{{
573-
ObjectMeta: metav1.ObjectMeta{Name: "gpu-template", Namespace: "ns"},
574-
Spec: resourcev1.ResourceClaimTemplateSpec{
575-
Spec: resourcev1.ResourceClaimSpec{
576-
Devices: resourcev1.DeviceClaim{
577-
Requests: []resourcev1.DeviceRequest{{
578-
Name: "gpu-request",
579-
Exactly: &resourcev1.ExactDeviceRequest{
580-
DeviceClassName: "unmapped.example.com",
581-
AllocationMode: resourcev1.DeviceAllocationModeExactCount,
582-
Count: 1,
583-
},
584-
}},
585-
},
586-
},
587-
},
588-
}},
539+
resourceClaimTemplates: []*resourcev1.ResourceClaimTemplate{
540+
utiltesting.MakeResourceClaimTemplate("gpu-template", "ns").
541+
DeviceRequest("gpu-request", "unmapped.example.com", 1).
542+
Obj(),
543+
},
589544
cq: utiltestingapi.MakeClusterQueue("cq").
590545
ResourceGroup(
591546
*utiltestingapi.MakeFlavorQuotas("flavor1").

‎pkg/dra/claims_test.go‎

Lines changed: 13 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ import (
3434
configapi "sigs.k8s.io/kueue/apis/config/v1beta1"
3535
kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
3636
"sigs.k8s.io/kueue/pkg/dra"
37+
utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
3738
)
3839

3940
func Test_GetResourceRequests(t *testing.T) {
@@ -42,37 +43,13 @@ func Test_GetResourceRequests(t *testing.T) {
4243
_ = kueue.AddToScheme(scheme)
4344
_ = resourcev1.AddToScheme(scheme)
4445

45-
tmpl := &resourcev1.ResourceClaimTemplate{
46-
ObjectMeta: metav1.ObjectMeta{Name: "claim-tmpl-1", Namespace: "ns1"},
47-
Spec: resourcev1.ResourceClaimTemplateSpec{
48-
Spec: resourcev1.ResourceClaimSpec{
49-
Devices: resourcev1.DeviceClaim{
50-
Requests: []resourcev1.DeviceRequest{{
51-
Exactly: &resourcev1.ExactDeviceRequest{
52-
AllocationMode: resourcev1.DeviceAllocationModeExactCount,
53-
Count: 2,
54-
DeviceClassName: "test-deviceclass-1",
55-
},
56-
}},
57-
},
58-
},
59-
},
60-
}
46+
tmpl := utiltesting.MakeResourceClaimTemplate("claim-tmpl-1", "ns1").
47+
DeviceRequest("device-request", "test-deviceclass-1", 2).
48+
Obj()
6149

62-
claim := &resourcev1.ResourceClaim{
63-
ObjectMeta: metav1.ObjectMeta{Name: "claim-2", Namespace: "ns1"},
64-
Spec: resourcev1.ResourceClaimSpec{
65-
Devices: resourcev1.DeviceClaim{
66-
Requests: []resourcev1.DeviceRequest{{
67-
Exactly: &resourcev1.ExactDeviceRequest{
68-
AllocationMode: resourcev1.DeviceAllocationModeExactCount,
69-
Count: 1,
70-
DeviceClassName: "test-deviceclass-2",
71-
},
72-
}},
73-
},
74-
},
75-
}
50+
claim := utiltesting.MakeResourceClaim("claim-2", "ns1").
51+
DeviceRequest("device-request", "test-deviceclass-2", 1).
52+
Obj()
7653

7754
wl := &kueue.Workload{
7855
ObjectMeta: metav1.ObjectMeta{Name: "wl", Namespace: "ns1"},
@@ -147,10 +124,9 @@ func Test_GetResourceRequests(t *testing.T) {
147124
}
148125
},
149126
extraObjects: []runtime.Object{
150-
&resourcev1.ResourceClaimTemplate{
151-
ObjectMeta: metav1.ObjectMeta{Name: "claim-tmpl-2", Namespace: "ns1"},
152-
Spec: resourcev1.ResourceClaimTemplateSpec{Spec: resourcev1.ResourceClaimSpec{Devices: resourcev1.DeviceClaim{Requests: []resourcev1.DeviceRequest{{Exactly: &resourcev1.ExactDeviceRequest{AllocationMode: resourcev1.DeviceAllocationModeExactCount, Count: 1, DeviceClassName: "test-deviceclass-2"}}}}}},
153-
},
127+
utiltesting.MakeResourceClaimTemplate("claim-tmpl-2", "ns1").
128+
DeviceRequest("device-request", "test-deviceclass-2", 1).
129+
Obj(),
154130
},
155131
lookup: func(dc corev1.ResourceName) (corev1.ResourceName, bool) {
156132
m := map[corev1.ResourceName]corev1.ResourceName{"test-deviceclass-1": "res-1", "test-deviceclass-2": "res-2"}
@@ -193,10 +169,9 @@ func Test_GetResourceRequests(t *testing.T) {
193169
{
194170
name: "Single template requesting two devices",
195171
extraObjects: []runtime.Object{
196-
&resourcev1.ResourceClaimTemplate{
197-
ObjectMeta: metav1.ObjectMeta{Name: "claim-tmpl-3", Namespace: "ns1"},
198-
Spec: resourcev1.ResourceClaimTemplateSpec{Spec: resourcev1.ResourceClaimSpec{Devices: resourcev1.DeviceClaim{Requests: []resourcev1.DeviceRequest{{Exactly: &resourcev1.ExactDeviceRequest{AllocationMode: resourcev1.DeviceAllocationModeExactCount, Count: 2, DeviceClassName: "test-deviceclass-1"}}}}}},
199-
},
172+
utiltesting.MakeResourceClaimTemplate("claim-tmpl-3", "ns1").
173+
DeviceRequest("device-request", "test-deviceclass-1", 2).
174+
Obj(),
200175
},
201176
modifyWL: func(w *kueue.Workload) {
202177
w.Spec.PodSets[0].Template.Spec.Containers = []corev1.Container{

0 commit comments

Comments
 (0)