From 65f3ad467f62619b732689b9a72ba6fdf95e4f66 Mon Sep 17 00:00:00 2001 From: Traian Schiau Date: Fri, 29 Sep 2023 17:10:19 +0300 Subject: [PATCH] [test/integration] Add provisioning ACC tests --- Makefile | 9 +- pkg/util/testing/wrappers.go | 19 + .../provisioning/provisioning_test.go | 383 ++++++++++++++++++ .../provisioning/suite_test.go | 73 ++++ test/util/util.go | 45 ++ 5 files changed, 528 insertions(+), 1 deletion(-) create mode 100644 test/integration/controller/admissionchecks/provisioning/provisioning_test.go create mode 100644 test/integration/controller/admissionchecks/provisioning/suite_test.go diff --git a/Makefile b/Makefile index a64ac393cac..6ec0f5818d8 100644 --- a/Makefile +++ b/Makefile @@ -153,7 +153,7 @@ test: generate gotestsum ## Run tests. $(GOTESTSUM) --junitfile $(ARTIFACTS)/junit.xml -- $(GO_TEST_FLAGS) $(shell $(GO_CMD) list ./... | grep -v '/test/') -coverpkg=./... -coverprofile $(ARTIFACTS)/cover.out .PHONY: test-integration -test-integration: manifests generate envtest ginkgo mpi-operator-crd ray-operator-crd jobset-operator-crd kf-training-operator-crd ## Run tests. +test-integration: manifests generate envtest ginkgo mpi-operator-crd ray-operator-crd jobset-operator-crd kf-training-operator-crd cluster-autoscaler-crd ## Run tests. KUBEBUILDER_ASSETS="$(shell $(ENVTEST) use $(ENVTEST_K8S_VERSION) -p path)" \ $(GINKGO) $(GINKGO_ARGS) --junit-report=junit.xml --output-dir=$(ARTIFACTS) -v $(INTEGRATION_TARGET) @@ -337,3 +337,10 @@ JOBSETROOT = $(shell $(GO_CMD) list -m -f "{{.Dir}}" sigs.k8s.io/jobset) jobset-operator-crd: mkdir -p $(PROJECT_DIR)/dep-crds/jobset-operator/ cp -f $(JOBSETROOT)/config/components/crd/bases/* $(PROJECT_DIR)/dep-crds/jobset-operator/ + + +CAROOT = $(shell $(GO_CMD) list -m -f "{{.Dir}}" k8s.io/autoscaler/cluster-autoscaler) +.PHONY: cluster-autoscaler-crd +cluster-autoscaler-crd: + mkdir -p $(PROJECT_DIR)/dep-crds/cluster-autoscaler/ + cp -f $(CAROOT)/config/crd/* $(PROJECT_DIR)/dep-crds/cluster-autoscaler/ diff --git a/pkg/util/testing/wrappers.go b/pkg/util/testing/wrappers.go index 9fa26f50f98..729a5a7729a 100644 --- a/pkg/util/testing/wrappers.go +++ b/pkg/util/testing/wrappers.go @@ -233,6 +233,11 @@ func (p *PodSetWrapper) Request(r corev1.ResourceName, q string) *PodSetWrapper return p } +func (p *PodSetWrapper) Iamage(image string) *PodSetWrapper { + p.Template.Spec.Containers[0].Image = image + return p +} + func (p *PodSetWrapper) SetMinimumCount(mc int32) *PodSetWrapper { p.MinCount = &mc return p @@ -578,6 +583,20 @@ func MakeAdmissionCheck(name string) *AdmissionCheckWrapper { } } +func (ac *AdmissionCheckWrapper) ControllerName(c string) *AdmissionCheckWrapper { + ac.Spec.ControllerName = c + return ac +} + +func (ac *AdmissionCheckWrapper) Parameters(apigroup, kind, name string) *AdmissionCheckWrapper { + ac.Spec.Parameters = &kueue.AdmissionCheckParametersReference{ + APIGroup: apigroup, + Kind: kind, + Name: name, + } + return ac +} + func (ac *AdmissionCheckWrapper) Obj() *kueue.AdmissionCheck { return &ac.AdmissionCheck } diff --git a/test/integration/controller/admissionchecks/provisioning/provisioning_test.go b/test/integration/controller/admissionchecks/provisioning/provisioning_test.go new file mode 100644 index 00000000000..b846843e84f --- /dev/null +++ b/test/integration/controller/admissionchecks/provisioning/provisioning_test.go @@ -0,0 +1,383 @@ +/* +Copyright 2023 The Kubernetes Authors. + +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 provisioning + +import ( + "github.com/google/go-cmp/cmp/cmpopts" + "github.com/onsi/ginkgo/v2" + "github.com/onsi/gomega" + corev1 "k8s.io/api/core/v1" + apimeta "k8s.io/apimachinery/pkg/api/meta" + "k8s.io/apimachinery/pkg/api/resource" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/types" + papi "k8s.io/autoscaler/cluster-autoscaler/provisioningrequest/apis/autoscaling.x-k8s.io/v1beta1" + "k8s.io/utils/ptr" + "sigs.k8s.io/controller-runtime/pkg/client" + + kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1" + "sigs.k8s.io/kueue/pkg/controller/admissionchecks/provisioning" + "sigs.k8s.io/kueue/pkg/util/testing" + "sigs.k8s.io/kueue/pkg/workload" + "sigs.k8s.io/kueue/test/util" +) + +var _ = ginkgo.Describe("Provisioning", func() { + + ginkgo.When("A workload is using a provision admission check", func() { + var ( + ns *corev1.Namespace + wlKey types.NamespacedName + ac *kueue.AdmissionCheck + prc *kueue.ProvisioningRequestConfig + rf *kueue.ResourceFlavor + admission *kueue.Admission + ) + ginkgo.BeforeEach(func() { + ns = &corev1.Namespace{ + ObjectMeta: metav1.ObjectMeta{ + GenerateName: "provisioning-", + }, + } + gomega.Expect(k8sClient.Create(ctx, ns)).To(gomega.Succeed()) + + wl := testing.MakeWorkload("wl", ns.Name). + PodSets( + *testing.MakePodSet("ps1", 3). + Request(corev1.ResourceCPU, "1"). + Iamage("iamge"). + Obj(), + *testing.MakePodSet("ps2", 6). + Request(corev1.ResourceCPU, "500m"). + Iamage("iamge"). + Obj(), + ). + Obj() + + gomega.Expect(k8sClient.Create(ctx, wl)).To(gomega.Succeed()) + + wlKey = client.ObjectKeyFromObject(wl) + + prc = &kueue.ProvisioningRequestConfig{ + ObjectMeta: metav1.ObjectMeta{ + Name: "prov-config", + }, + Spec: kueue.ProvisioningRequestConfigSpec{ + ProvisioningClassName: "provisioning-class", + Parameters: map[string]kueue.Parameter{ + "p1": "v1", + "p2": "v2", + }, + }, + } + gomega.Expect(k8sClient.Create(ctx, prc)).To(gomega.Succeed()) + + ac = testing.MakeAdmissionCheck("ac-prov"). + ControllerName(provisioning.ControllerName). + Parameters(kueue.GroupVersion.Group, "ProvisioningRequestConfig", prc.Name). + Obj() + gomega.Expect(k8sClient.Create(ctx, ac)).To(gomega.Succeed()) + + rf = testing.MakeResourceFlavor("rf1").Label("ns1", "ns1v").Obj() + gomega.Expect(k8sClient.Create(ctx, rf)).To(gomega.Succeed()) + + admission = testing.MakeAdmission("q"). + PodSets( + kueue.PodSetAssignment{ + Name: "ps1", + Flavors: map[corev1.ResourceName]kueue.ResourceFlavorReference{ + corev1.ResourceCPU: kueue.ResourceFlavorReference(rf.Name), + }, + ResourceUsage: map[corev1.ResourceName]resource.Quantity{ + corev1.ResourceCPU: resource.MustParse("3"), + }, + Count: ptr.To[int32](3), + }, + kueue.PodSetAssignment{ + Name: "ps2", + Flavors: map[corev1.ResourceName]kueue.ResourceFlavorReference{ + corev1.ResourceCPU: kueue.ResourceFlavorReference(rf.Name), + }, + ResourceUsage: map[corev1.ResourceName]resource.Quantity{ + corev1.ResourceCPU: resource.MustParse("2"), + }, + Count: ptr.To[int32](4), + }, + ). + Obj() + + }) + + ginkgo.AfterEach(func() { + util.ExpectResourceFlavorToBeDeleted(ctx, k8sClient, rf, true) + util.ExpectAdmissionCheckToBeDeleted(ctx, k8sClient, ac, true) + util.ExpectProvisioningRequestConfigToBeDeleted(ctx, k8sClient, prc, true) + gomega.Expect(util.DeleteNamespace(ctx, k8sClient, ns)).To(gomega.Succeed()) + }) + + ginkgo.It("Should not create provisioning requests before quota is reserved", func() { + ginkgo.By("Setting the admission check to the workload", func() { + uodatedWl := &kueue.Workload{} + gomega.Eventually(func() error { + err := k8sClient.Get(ctx, wlKey, uodatedWl) + if err != nil { + return err + } + util.SetWorkloadsAdmissionCkeck(ctx, k8sClient, uodatedWl, ac.Name, kueue.CheckStatePending, false) + return nil + }, util.Timeout, util.Interval).Should(gomega.Succeed()) + }) + + ginkgo.By("Checking no provision request is created", func() { + provReqKey := types.NamespacedName{ + Namespace: wlKey.Namespace, + Name: provisioning.GetProvisioningRequestName(wlKey.Name, ac.Name), + } + gomega.Consistently(func() error { + request := &papi.ProvisioningRequest{} + return k8sClient.Get(ctx, provReqKey, request) + }, util.ConsistentDuration, util.Interval).Should(testing.BeNotFoundError()) + }) + }) + + ginkgo.It("Should create provisioning requests after quota is reserved and remove it when reservation is lost", func() { + uodatedWl := &kueue.Workload{} + ginkgo.By("Setting the admission check to the workload", func() { + gomega.Eventually(func() error { + err := k8sClient.Get(ctx, wlKey, uodatedWl) + if err != nil { + return err + } + util.SetWorkloadsAdmissionCkeck(ctx, k8sClient, uodatedWl, ac.Name, kueue.CheckStatePending, false) + return nil + }, util.Timeout, util.Interval).Should(gomega.Succeed()) + }) + + ginkgo.By("Setting the quota reservation to the workload", func() { + uodatedWl := &kueue.Workload{} + gomega.Eventually(func() error { + err := k8sClient.Get(ctx, wlKey, uodatedWl) + if err != nil { + return err + } + gomega.Expect(util.SetQuotaReservation(ctx, k8sClient, uodatedWl, admission)).To(gomega.Succeed()) + util.SetWorkloadsAdmissionCkeck(ctx, k8sClient, uodatedWl, ac.Name, kueue.CheckStatePending, false) + return nil + }, util.Timeout, util.Interval).Should(gomega.Succeed()) + }) + + createdRequest := &papi.ProvisioningRequest{} + ginkgo.By("Checking that the provision request is created", func() { + provReqKey := types.NamespacedName{ + Namespace: wlKey.Namespace, + Name: provisioning.GetProvisioningRequestName(wlKey.Name, ac.Name), + } + gomega.Eventually(func() error { + return k8sClient.Get(ctx, provReqKey, createdRequest) + }, util.Timeout, util.Interval).Should(gomega.Succeed()) + }) + + ignoreContainersDefaults := cmpopts.IgnoreFields(corev1.Container{}, "TerminationMessagePath", "TerminationMessagePolicy", "ImagePullPolicy") + ginkgo.By("Checking that the provision requests content", func() { + gomega.Expect(createdRequest.Spec.ProvisioningClassName).To(gomega.Equal("provisioning-class")) + gomega.Expect(createdRequest.Spec.Parameters).To(gomega.BeComparableTo(map[string]papi.Parameter{ + "p1": "v1", + "p2": "v2", + })) + gomega.Expect(createdRequest.Spec.PodSets).To(gomega.HaveLen(2)) + + ps1 := createdRequest.Spec.PodSets[0] + gomega.Expect(ps1.Count).To(gomega.Equal(int32(3))) + gomega.Expect(ps1.PodTemplateRef.Name).NotTo(gomega.BeEmpty()) + + // check the created pod template + createdTemplate := &corev1.PodTemplate{} + templateKey := types.NamespacedName{ + Namespace: createdRequest.Namespace, + Name: ps1.PodTemplateRef.Name, + } + gomega.Eventually(func() error { + return k8sClient.Get(ctx, templateKey, createdTemplate) + }, util.Timeout, util.Interval).Should(gomega.Succeed()) + gomega.Expect(createdTemplate.Template.Spec.Containers).To(gomega.BeComparableTo(uodatedWl.Spec.PodSets[0].Template.Spec.Containers, ignoreContainersDefaults)) + gomega.Expect(createdTemplate.Template.Spec.NodeSelector).To(gomega.BeComparableTo(map[string]string{"ns1": "ns1v"})) + + ps2 := createdRequest.Spec.PodSets[1] + gomega.Expect(ps2.Count).To(gomega.Equal(int32(4))) + gomega.Expect(ps2.PodTemplateRef.Name).NotTo(gomega.BeEmpty()) + + // check the created pod template + templateKey.Name = ps2.PodTemplateRef.Name + gomega.Eventually(func() error { + return k8sClient.Get(ctx, templateKey, createdTemplate) + }, util.Timeout, util.Interval).Should(gomega.Succeed()) + gomega.Expect(createdTemplate.Template.Spec.Containers).To(gomega.BeComparableTo(uodatedWl.Spec.PodSets[1].Template.Spec.Containers, ignoreContainersDefaults)) + gomega.Expect(createdTemplate.Template.Spec.NodeSelector).To(gomega.BeComparableTo(map[string]string{"ns1": "ns1v"})) + }) + + ginkgo.By("Removing the quota reservation from the workload", func() { + uodatedWl := &kueue.Workload{} + gomega.Eventually(func() error { + err := k8sClient.Get(ctx, wlKey, uodatedWl) + if err != nil { + return err + } + gomega.Expect(util.SetQuotaReservation(ctx, k8sClient, uodatedWl, nil)).To(gomega.Succeed()) + util.SetWorkloadsAdmissionCkeck(ctx, k8sClient, uodatedWl, ac.Name, kueue.CheckStatePending, false) + return nil + }, util.Timeout, util.Interval).Should(gomega.Succeed()) + }) + + ginkgo.By("Checking no provision request is deleted", func() { + provReqKey := types.NamespacedName{ + Namespace: wlKey.Namespace, + Name: provisioning.GetProvisioningRequestName(wlKey.Name, ac.Name), + } + gomega.Eventually(func() error { + request := &papi.ProvisioningRequest{} + return k8sClient.Get(ctx, provReqKey, request) + }, util.Timeout, util.Interval).Should(testing.BeNotFoundError()) + }) + }) + + ginkgo.It("Should set the condition ready when the provision succeed", func() { + ginkgo.By("Setting the admission check to the workload", func() { + uodatedWl := &kueue.Workload{} + gomega.Eventually(func() error { + err := k8sClient.Get(ctx, wlKey, uodatedWl) + if err != nil { + return err + } + util.SetWorkloadsAdmissionCkeck(ctx, k8sClient, uodatedWl, ac.Name, kueue.CheckStatePending, false) + return nil + }, util.Timeout, util.Interval).Should(gomega.Succeed()) + }) + + ginkgo.By("Setting the quota reservation to the workload", func() { + uodatedWl := &kueue.Workload{} + gomega.Eventually(func() error { + err := k8sClient.Get(ctx, wlKey, uodatedWl) + if err != nil { + return err + } + gomega.Expect(util.SetQuotaReservation(ctx, k8sClient, uodatedWl, admission)).To(gomega.Succeed()) + util.SetWorkloadsAdmissionCkeck(ctx, k8sClient, uodatedWl, ac.Name, kueue.CheckStatePending, false) + return nil + }, util.Timeout, util.Interval).Should(gomega.Succeed()) + }) + + provReqKey := types.NamespacedName{ + Namespace: wlKey.Namespace, + Name: provisioning.GetProvisioningRequestName(wlKey.Name, ac.Name), + } + ginkgo.By("Setting the provision request as Provisioned", func() { + createdRequest := &papi.ProvisioningRequest{} + gomega.Eventually(func() error { + err := k8sClient.Get(ctx, provReqKey, createdRequest) + if err != nil { + return err + } + apimeta.SetStatusCondition(&createdRequest.Status.Conditions, metav1.Condition{ + Type: papi.Provisioned, + Status: metav1.ConditionTrue, + Reason: papi.Provisioned, + }) + return k8sClient.Status().Update(ctx, createdRequest) + }, util.Timeout, util.Interval).Should(gomega.Succeed()) + }) + + ginkgo.By("Checking the admission check", func() { + uodatedWl := &kueue.Workload{} + gomega.Eventually(func(g gomega.Gomega) { + g.Expect(k8sClient.Get(ctx, wlKey, uodatedWl)).To(gomega.Succeed()) + state := workload.FindAdmissionCheck(uodatedWl.Status.AdmissionChecks, ac.Name) + g.Expect(state).NotTo(gomega.BeNil()) + g.Expect(state.State).To(gomega.Equal(kueue.CheckStateReady)) + g.Expect(state.PodSetUpdates).To(gomega.BeComparableTo([]kueue.PodSetUpdate{ + { + Name: "ps1", + Annotations: map[string]string{ + provisioning.ConsumesAnnotationKey: provReqKey.Name, + }, + }, + { + Name: "ps2", + Annotations: map[string]string{ + provisioning.ConsumesAnnotationKey: provReqKey.Name, + }, + }, + })) + }, util.Timeout, util.Interval).Should(gomega.Succeed()) + }) + }) + ginkgo.It("Should set the condition rejected when the provision fails", func() { + ginkgo.By("Setting the admission check to the workload", func() { + uodatedWl := &kueue.Workload{} + gomega.Eventually(func() error { + err := k8sClient.Get(ctx, wlKey, uodatedWl) + if err != nil { + return err + } + util.SetWorkloadsAdmissionCkeck(ctx, k8sClient, uodatedWl, ac.Name, kueue.CheckStatePending, false) + return nil + }, util.Timeout, util.Interval).Should(gomega.Succeed()) + }) + + ginkgo.By("Setting the quota reservation to the workload", func() { + uodatedWl := &kueue.Workload{} + gomega.Eventually(func() error { + err := k8sClient.Get(ctx, wlKey, uodatedWl) + if err != nil { + return err + } + gomega.Expect(util.SetQuotaReservation(ctx, k8sClient, uodatedWl, admission)).To(gomega.Succeed()) + util.SetWorkloadsAdmissionCkeck(ctx, k8sClient, uodatedWl, ac.Name, kueue.CheckStatePending, false) + return nil + }, util.Timeout, util.Interval).Should(gomega.Succeed()) + }) + + ginkgo.By("Setting the provision request as Provisioned", func() { + createdRequest := &papi.ProvisioningRequest{} + provReqKey := types.NamespacedName{ + Namespace: wlKey.Namespace, + Name: provisioning.GetProvisioningRequestName(wlKey.Name, ac.Name), + } + gomega.Eventually(func() error { + err := k8sClient.Get(ctx, provReqKey, createdRequest) + if err != nil { + return err + } + apimeta.SetStatusCondition(&createdRequest.Status.Conditions, metav1.Condition{ + Type: papi.Failed, + Status: metav1.ConditionTrue, + Reason: papi.Failed, + }) + return k8sClient.Status().Update(ctx, createdRequest) + }, util.Timeout, util.Interval).Should(gomega.Succeed()) + }) + + ginkgo.By("Checking the admission check", func() { + uodatedWl := &kueue.Workload{} + gomega.Eventually(func(g gomega.Gomega) { + g.Expect(k8sClient.Get(ctx, wlKey, uodatedWl)).To(gomega.Succeed()) + state := workload.FindAdmissionCheck(uodatedWl.Status.AdmissionChecks, ac.Name) + g.Expect(state).NotTo(gomega.BeNil()) + g.Expect(state.State).To(gomega.Equal(kueue.CheckStateRejected)) + }, util.Timeout, util.Interval).Should(gomega.Succeed()) + }) + }) + }) +}) diff --git a/test/integration/controller/admissionchecks/provisioning/suite_test.go b/test/integration/controller/admissionchecks/provisioning/suite_test.go new file mode 100644 index 00000000000..f2d4cc2356b --- /dev/null +++ b/test/integration/controller/admissionchecks/provisioning/suite_test.go @@ -0,0 +1,73 @@ +/* +Copyright 2023 The Kubernetes Authors. + +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 provisioning + +import ( + "context" + "path/filepath" + "testing" + + "github.com/onsi/ginkgo/v2" + "github.com/onsi/gomega" + "k8s.io/client-go/rest" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/manager" + + "sigs.k8s.io/kueue/pkg/controller/admissionchecks/provisioning" + "sigs.k8s.io/kueue/test/integration/framework" + //+kubebuilder:scaffold:imports +) + +var ( + cfg *rest.Config + k8sClient client.Client + ctx context.Context + fwk *framework.Framework +) + +func TestProvisioning(t *testing.T) { + gomega.RegisterFailHandler(ginkgo.Fail) + + ginkgo.RunSpecs(t, + "Provisioning admission check suite", + ) +} + +func managerSetup(mgr manager.Manager, ctx context.Context) { + err := provisioning.AddToScheme(mgr.GetScheme()) + gomega.Expect(err).NotTo(gomega.HaveOccurred()) + + err = provisioning.SetupIndexer(ctx, mgr.GetFieldIndexer()) + gomega.Expect(err).NotTo(gomega.HaveOccurred()) + + reconciler := provisioning.NewController(mgr.GetClient()) + err = reconciler.SetupWithManager(mgr) + gomega.Expect(err).NotTo(gomega.HaveOccurred()) +} + +var _ = ginkgo.BeforeSuite(func() { + fwk = &framework.Framework{ + CRDPath: filepath.Join("..", "..", "..", "..", "..", "config", "components", "crd", "bases"), + DepCRDPaths: []string{filepath.Join("..", "..", "..", "..", "..", "dep-crds", "cluster-autoscaler")}, + } + cfg = fwk.Init() + ctx, k8sClient = fwk.RunManager(cfg, managerSetup) +}) + +var _ = ginkgo.AfterSuite(func() { + fwk.Teardown() +}) diff --git a/test/util/util.go b/test/util/util.go index 28906cffec1..4843db1e782 100644 --- a/test/util/util.go +++ b/test/util/util.go @@ -45,6 +45,15 @@ func DeleteAdmissionCheck(ctx context.Context, c client.Client, ac *kueue.Admiss return nil } +func DeleteProvisioningRequestConfig(ctx context.Context, c client.Client, ac *kueue.ProvisioningRequestConfig) error { + if ac != nil { + if err := c.Delete(ctx, ac); err != nil && !apierrors.IsNotFound(err) { + return err + } + } + return nil +} + func DeleteWorkload(ctx context.Context, c client.Client, wl *kueue.Workload) error { if wl != nil { if err := c.Delete(ctx, wl); err != nil && !apierrors.IsNotFound(err) { @@ -299,6 +308,19 @@ func ExpectAdmissionCheckToBeDeleted(ctx context.Context, k8sClient client.Clien }, Timeout, Interval).Should(testing.BeNotFoundError()) } +func ExpectProvisioningRequestConfigToBeDeleted(ctx context.Context, k8sClient client.Client, prc *kueue.ProvisioningRequestConfig, deleteAC bool) { + if prc == nil { + return + } + if deleteAC { + gomega.Expect(client.IgnoreNotFound(DeleteProvisioningRequestConfig(ctx, k8sClient, prc))).To(gomega.Succeed()) + } + gomega.EventuallyWithOffset(1, func() error { + var newAC kueue.AdmissionCheck + return k8sClient.Get(ctx, client.ObjectKeyFromObject(prc), &newAC) + }, Timeout, Interval).Should(testing.BeNotFoundError()) +} + func ExpectClusterQueueToBeDeleted(ctx context.Context, k8sClient client.Client, cq *kueue.ClusterQueue, deleteCq bool) { if deleteCq { gomega.Expect(DeleteClusterQueue(ctx, k8sClient, cq)).ToNot(gomega.HaveOccurred()) @@ -402,3 +424,26 @@ func FinishEvictionForWorkloads(ctx context.Context, k8sClient client.Client, wl } } + +func SetWorkloadsAdmissionCkeck(ctx context.Context, k8sClient client.Client, wl *kueue.Workload, check string, state kueue.CheckState, expectExisting bool) { + var updatedWorkload kueue.Workload + + gomega.EventuallyWithOffset(1, func() error { + err := k8sClient.Get(ctx, client.ObjectKeyFromObject(wl), &updatedWorkload) + if err != nil { + return err + } + var checkCondition *kueue.AdmissionCheckState + if expectExisting { + checkState := workload.FindAdmissionCheck(updatedWorkload.Status.AdmissionChecks, check) + gomega.ExpectWithOffset(1, checkState).NotTo(gomega.BeNil(), "the check was not found") + } else { + checkCondition = &kueue.AdmissionCheckState{ + Name: check, + } + } + checkCondition.State = state + workload.SetAdmissionCheckState(&updatedWorkload.Status.AdmissionChecks, *checkCondition) + return k8sClient.Status().Update(ctx, &updatedWorkload) + }, Timeout, Interval).Should(gomega.Succeed()) +}