Skip to content

Commit

Permalink
Add TestReconcile unit tests
Browse files Browse the repository at this point in the history
  • Loading branch information
PBundyra committed Apr 10, 2024
1 parent bbcd05d commit 7c77a63
Show file tree
Hide file tree
Showing 6 changed files with 140 additions and 73 deletions.
12 changes: 7 additions & 5 deletions charts/kueue/templates/crd/kueue.x-k8s.io_clusterqueues.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -74,8 +74,9 @@ spec:
description: ClusterQueueSpec defines the desired state of ClusterQueue
properties:
admissionCheckStrategy:
description: admissionCheckStrategy is the list of AdmissionChecks
that should run when a particular ResourceFlavor is used.
description: |-
admissionCheckStrategy is the list of AdmissionChecks that should run when a particular ResourceFlavor is used.
Cannot be used along with AdmissionChecks.
items:
properties:
name:
Expand All @@ -84,7 +85,7 @@ spec:
onFlavors:
description: |-
forFlavors is a list of ResourceFlavors' metav1.Names that this AdmissionCheck should run for.
if empty the AdmissionCheck will run for all workloads submitted to the ClusterQueue.
If empty, the AdmissionCheck will run for all workloads submitted to the ClusterQueue.
items:
type: string
type: array
Expand All @@ -94,8 +95,9 @@ spec:
type: object
type: array
admissionChecks:
description: admissionChecks lists the AdmissionChecks required by
this ClusterQueue
description: |-
admissionChecks lists the AdmissionChecks required by this ClusterQueue.
Cannot be used along with AdmissionCheckStrategy.
items:
type: string
type: array
Expand Down
12 changes: 7 additions & 5 deletions config/components/crd/bases/kueue.x-k8s.io_clusterqueues.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -59,8 +59,9 @@ spec:
description: ClusterQueueSpec defines the desired state of ClusterQueue
properties:
admissionCheckStrategy:
description: admissionCheckStrategy is the list of AdmissionChecks
that should run when a particular ResourceFlavor is used.
description: |-
admissionCheckStrategy is the list of AdmissionChecks that should run when a particular ResourceFlavor is used.
Cannot be used along with AdmissionChecks.
items:
properties:
name:
Expand All @@ -69,7 +70,7 @@ spec:
onFlavors:
description: |-
forFlavors is a list of ResourceFlavors' metav1.Names that this AdmissionCheck should run for.
if empty the AdmissionCheck will run for all workloads submitted to the ClusterQueue.
If empty, the AdmissionCheck will run for all workloads submitted to the ClusterQueue.
items:
type: string
type: array
Expand All @@ -79,8 +80,9 @@ spec:
type: object
type: array
admissionChecks:
description: admissionChecks lists the AdmissionChecks required by
this ClusterQueue
description: |-
admissionChecks lists the AdmissionChecks required by this ClusterQueue.
Cannot be used along with AdmissionCheckStrategy.
items:
type: string
type: array
Expand Down
124 changes: 63 additions & 61 deletions pkg/controller/core/workload_controller_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,6 @@ import (
"k8s.io/apimachinery/pkg/types"
testingclock "k8s.io/utils/clock/testing"
"k8s.io/utils/ptr"
ctrl "sigs.k8s.io/controller-runtime"
"sigs.k8s.io/controller-runtime/pkg/client"
"sigs.k8s.io/controller-runtime/pkg/client/interceptor"
"sigs.k8s.io/controller-runtime/pkg/reconcile"
Expand Down Expand Up @@ -310,6 +309,7 @@ var (
kueue.Workload{}, "TypeMeta", "ObjectMeta.ResourceVersion", "Status.RequeueState.RequeueAt",
),
cmpopts.IgnoreFields(metav1.Condition{}, "LastTransitionTime"),
cmpopts.IgnoreFields(kueue.AdmissionCheckState{}, "LastTransitionTime"),
cmpopts.SortSlices(func(a, b metav1.Condition) bool { return a.Type < b.Type }),
}
)
Expand All @@ -318,11 +318,61 @@ func TestReconcile(t *testing.T) {
testStartTime := time.Now()
cases := map[string]struct {
workload *kueue.Workload
cq *kueue.ClusterQueue
lq *kueue.LocalQueue
wantWorkload *kueue.Workload
wantError error
wantEvents []utiltesting.EventRecord
reconcilerOpts []Option
}{
"assign Admission Checks from ClusterQueue.spec.AdmissionCheckStrategy": {
workload: utiltesting.MakeWorkload("wl", "ns").
ReserveQuota(utiltesting.MakeAdmission("cq").Assignment("cpu", "flavor1", "1").Obj()).
Queue("queue").
Obj(),
cq: utiltesting.MakeClusterQueue("cq").
AdmissionCheckStrategy(
*utiltesting.MakeAdmissionCheckStrategy("ac1", "flavor1").Obj(),
*utiltesting.MakeAdmissionCheckStrategy("ac2").Obj()).
Obj(),
lq: utiltesting.MakeLocalQueue("queue", "ns").ClusterQueue("cq").Obj(),
wantWorkload: utiltesting.MakeWorkload("wl", "ns").
ReserveQuota(utiltesting.MakeAdmission("cq").Assignment("cpu", "flavor1", "1").Obj()).
Queue("queue").
AdmissionChecks(
kueue.AdmissionCheckState{
Name: "ac1",
State: kueue.CheckStatePending,
},
kueue.AdmissionCheckState{
Name: "ac2",
State: kueue.CheckStatePending,
}).
Obj(),
},
"assign Admission Checks from ClusterQueue.spec.AdmissionChecks": {
workload: utiltesting.MakeWorkload("wl", "ns").
ReserveQuota(utiltesting.MakeAdmission("cq").Assignment("cpu", "flavor1", "1").Obj()).
Queue("queue").
Obj(),
cq: utiltesting.MakeClusterQueue("cq").
AdmissionChecks("ac1", "ac2").
Obj(),
lq: utiltesting.MakeLocalQueue("queue", "ns").ClusterQueue("cq").Obj(),
wantWorkload: utiltesting.MakeWorkload("wl", "ns").
ReserveQuota(utiltesting.MakeAdmission("cq").Assignment("cpu", "flavor1", "1").Obj()).
Queue("queue").
AdmissionChecks(
kueue.AdmissionCheckState{
Name: "ac1",
State: kueue.CheckStatePending,
},
kueue.AdmissionCheckState{
Name: "ac2",
State: kueue.CheckStatePending,
}).
Obj(),
},
"admit": {
workload: utiltesting.MakeWorkload("wl", "ns").
ReserveQuota(utiltesting.MakeAdmission("q1").Obj()).
Expand Down Expand Up @@ -545,6 +595,18 @@ func TestReconcile(t *testing.T) {
ctx, ctxCancel := context.WithCancel(ctxWithLogger)
defer ctxCancel()

if tc.cq != nil {
if err := cl.Create(ctx, tc.cq); err != nil {
t.Error("couldn't create the cluster queue", err)
}
if err := qManager.AddClusterQueue(ctx, tc.cq); err != nil {
t.Error("couldn't add the cluster queue to the cache", err)
}
if err := qManager.AddLocalQueue(ctx, tc.lq); err != nil {
t.Error("couldn't add the local queue to the cache", err)
}
}

_, gotError := reconciler.Reconcile(ctx, reconcile.Request{NamespacedName: client.ObjectKeyFromObject(tc.workload)})

if diff := cmp.Diff(tc.wantError, gotError); diff != "" {
Expand Down Expand Up @@ -586,63 +648,3 @@ func TestReconcile(t *testing.T) {
})
}
}

func TestAdmissionCheckStrategy(t *testing.T) {
cases := map[string]struct {
cq *kueue.ClusterQueue
wl *kueue.Workload
wantAdmissionChecks []string
}{
"AdmissionCheckStrategy with a flavor": {
wl: utiltesting.MakeWorkload("wl", "ns").
ReserveQuota(utiltesting.MakeAdmission("cq").Assignment("cpu", "flavor1", "1").Obj()).
Obj(),
cq: utiltesting.MakeClusterQueue("cq").
AdmissionCheckStrategy(*utiltesting.MakeAdmissionCheckStrategy("ac1", "flavor1").Obj()).
Obj(),
wantAdmissionChecks: []string{"ac1"},
},
"AdmissionCheckStrategy without a flavor": {
wl: utiltesting.MakeWorkload("wl", "ns").
ReserveQuota(utiltesting.MakeAdmission("cq").Assignment("cpu", "flavor1", "1").Obj()).
Obj(),
cq: utiltesting.MakeClusterQueue("cq").
AdmissionCheckStrategy(*utiltesting.MakeAdmissionCheckStrategy("ac1").Obj()).
Obj(),
wantAdmissionChecks: []string{"ac1"},
},
"Two AdmissionCheckStrategies, one with flavor, one without flavor": {
wl: utiltesting.MakeWorkload("wl", "ns").
ReserveQuota(utiltesting.MakeAdmission("cq").Assignment("cpu", "flavor1", "1").Obj()).
Obj(),
cq: utiltesting.MakeClusterQueue("cq").
AdmissionCheckStrategy(
*utiltesting.MakeAdmissionCheckStrategy("ac1", "flavor1").Obj(),
*utiltesting.MakeAdmissionCheckStrategy("ac2").Obj()).
Obj(),
wantAdmissionChecks: []string{"ac1", "ac2"},
},
"Workload has no Admission": {
wl: utiltesting.MakeWorkload("wl", "ns").
Obj(),
cq: utiltesting.MakeClusterQueue("cq").
AdmissionCheckStrategy(
*utiltesting.MakeAdmissionCheckStrategy("ac1", "flavor1").Obj(),
*utiltesting.MakeAdmissionCheckStrategy("ac2").Obj()).
Obj(),
wantAdmissionChecks: []string{},
},
}
for name, tc := range cases {
t.Run(name, func(t *testing.T) {
ctx := context.Background()
log := ctrl.LoggerFrom(ctx)
gotAdmissionChecks := getChecksFromACStrategies(&log, tc.wl, tc.cq)

if diff := cmp.Diff(tc.wantAdmissionChecks, gotAdmissionChecks); diff != "" {
t.Errorf("Unexpected AdmissionChecks, (want-/got+): %s", diff)
}
})

}
}
2 changes: 1 addition & 1 deletion pkg/workload/workload.go
Original file line number Diff line number Diff line change
Expand Up @@ -556,4 +556,4 @@ func AdmissionChecksForWorkload(log *logr.Logger, wl *kueue.Workload, cq *kueue.
}
}
return admissionCheckNames
}
}
61 changes: 61 additions & 0 deletions pkg/workload/workload_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ import (
"k8s.io/apimachinery/pkg/api/resource"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/utils/ptr"
ctrl "sigs.k8s.io/controller-runtime"
"sigs.k8s.io/controller-runtime/pkg/client"

config "sigs.k8s.io/kueue/apis/config/v1beta1"
Expand Down Expand Up @@ -643,3 +644,63 @@ func TestResourceUsage(t *testing.T) {
})
}
}

func TestAdmissionCheckStrategy(t *testing.T) {
cases := map[string]struct {
cq *kueue.ClusterQueue
wl *kueue.Workload
wantAdmissionChecks []string
}{
"AdmissionCheckStrategy with a flavor": {
wl: utiltesting.MakeWorkload("wl", "ns").
ReserveQuota(utiltesting.MakeAdmission("cq").Assignment("cpu", "flavor1", "1").Obj()).
Obj(),
cq: utiltesting.MakeClusterQueue("cq").
AdmissionCheckStrategy(*utiltesting.MakeAdmissionCheckStrategy("ac1", "flavor1").Obj()).
Obj(),
wantAdmissionChecks: []string{"ac1"},
},
"AdmissionCheckStrategy without a flavor": {
wl: utiltesting.MakeWorkload("wl", "ns").
ReserveQuota(utiltesting.MakeAdmission("cq").Assignment("cpu", "flavor1", "1").Obj()).
Obj(),
cq: utiltesting.MakeClusterQueue("cq").
AdmissionCheckStrategy(*utiltesting.MakeAdmissionCheckStrategy("ac1").Obj()).
Obj(),
wantAdmissionChecks: []string{"ac1"},
},
"Two AdmissionCheckStrategies, one with flavor, one without flavor": {
wl: utiltesting.MakeWorkload("wl", "ns").
ReserveQuota(utiltesting.MakeAdmission("cq").Assignment("cpu", "flavor1", "1").Obj()).
Obj(),
cq: utiltesting.MakeClusterQueue("cq").
AdmissionCheckStrategy(
*utiltesting.MakeAdmissionCheckStrategy("ac1", "flavor1").Obj(),
*utiltesting.MakeAdmissionCheckStrategy("ac2").Obj()).
Obj(),
wantAdmissionChecks: []string{"ac1", "ac2"},
},
"Workload has no Admission": {
wl: utiltesting.MakeWorkload("wl", "ns").
Obj(),
cq: utiltesting.MakeClusterQueue("cq").
AdmissionCheckStrategy(
*utiltesting.MakeAdmissionCheckStrategy("ac1", "flavor1").Obj(),
*utiltesting.MakeAdmissionCheckStrategy("ac2").Obj()).
Obj(),
wantAdmissionChecks: nil,
},
}
for name, tc := range cases {
t.Run(name, func(t *testing.T) {
ctx := context.Background()
log := ctrl.LoggerFrom(ctx)
gotAdmissionChecks := AdmissionChecksForWorkload(&log, tc.wl, tc.cq)

if diff := cmp.Diff(tc.wantAdmissionChecks, gotAdmissionChecks); diff != "" {
t.Errorf("Unexpected AdmissionChecks, (want-/got+): %s", diff)
}
})

}
}
2 changes: 1 addition & 1 deletion test/integration/scheduler/workload_controller_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ var _ = ginkgo.Describe("Workload controller with scheduler", func() {
check2 *kueue.AdmissionCheck
check3 *kueue.AdmissionCheck
reservationFlavor string = "reservation"
updatedWl kueue.Workload
updatedWl kueue.Workload
flavorOnDemand string = "on-demand"
resourceGPU corev1.ResourceName = "example.com/gpu"
)
Expand Down

0 comments on commit 7c77a63

Please sign in to comment.