Repository: kubernetes-sigs/kueue

Issue:
Extract iteration over findFlavorForPodSetResource to a new helper function
<!-- Please only use this template for submitting clean up requests -->

**What would you like to be cleaned**:

The findFlavorForPodSetResource is long and the new iteration over PodSets increases indentention and requires helper var flavorUnmatchMessage

**Why is this needed**:

The function is getting hard to read with many nested loops.

See review comment: https://github.com/kubernetes-sigs/kueue/pull/5878#discussion_r2216104567

Discussion:
@olekzabl: /assign

Solution patch:
diff --git a/Makefile-test.mk b/Makefile-test.mk
index 7c89fc764..173536019 100644
--- a/Makefile-test.mk
+++ b/Makefile-test.mk
@@ -51,6 +51,9 @@ E2E_USE_HELM ?= false
 # Default will delete default kind cluster
 KIND_CLUSTER_NAME ?= kind
 
+# For restricting to a specific directory
+GO_TEST_TARGET ?= .
+
 ##@ Tests
 
 # Periodic builds are tested with full ray image
@@ -62,7 +65,7 @@ endif
 
 .PHONY: test
 test: gotestsum ## Run tests.
-	$(GOTESTSUM) --junitfile $(ARTIFACTS)/junit.xml -- $(GOFLAGS) $(GO_TEST_FLAGS) $(shell $(GO_CMD) list ./... | grep -v '/test/') -coverpkg=./... -coverprofile $(ARTIFACTS)/cover.out
+	$(GOTESTSUM) --junitfile $(ARTIFACTS)/junit.xml -- $(GOFLAGS) $(GO_TEST_FLAGS) $(shell $(GO_CMD) list $(GO_TEST_TARGET)/... | grep -v '/test/') -coverpkg=$(GO_TEST_TARGET)/... -coverprofile $(ARTIFACTS)/cover.out
 
 .PHONY: test-integration
 test-integration: gomod-download envtest ginkgo dep-crds kueuectl ginkgo-top ## Run integration tests for all singlecluster suites.
diff --git a/pkg/controller/admissionchecks/multikueue/workload.go b/pkg/controller/admissionchecks/multikueue/workload.go
index 71f3245dd..d806ea040 100644
--- a/pkg/controller/admissionchecks/multikueue/workload.go
+++ b/pkg/controller/admissionchecks/multikueue/workload.go
@@ -282,7 +282,7 @@ func (w *wlReconciler) multikueueAC(ctx context.Context, local *kueue.Workload)
 	if len(relevantChecks) == 0 {
 		return nil, nil
 	}
-	return workload.FindAdmissionCheck(local.Status.AdmissionChecks, relevantChecks[0]), nil
+	return admissioncheck.FindAdmissionCheck(local.Status.AdmissionChecks, relevantChecks[0]), nil
 }
 
 func (w *wlReconciler) adapter(local *kueue.Workload) (jobframework.MultiKueueAdapter, *metav1.OwnerReference) {
@@ -331,7 +331,7 @@ func (w *wlReconciler) reconcileGroup(ctx context.Context, group *wlGroup) (reco
 	log := ctrl.LoggerFrom(ctx).WithValues("op", "reconcileGroup")
 	log.V(3).Info("Reconcile Workload Group")
 
-	acs := workload.FindAdmissionCheck(group.local.Status.AdmissionChecks, group.acName)
+	acs := admissioncheck.FindAdmissionCheck(group.local.Status.AdmissionChecks, group.acName)
 
 	// 1. delete all remote workloads when finished or the local wl has no reservation
 	if group.IsFinished() || !workload.HasQuotaReservation(group.local) {
@@ -395,7 +395,7 @@ func (w *wlReconciler) reconcileGroup(ctx context.Context, group *wlGroup) (reco
 			}
 		}
 
-		acs := workload.FindAdmissionCheck(group.local.Status.AdmissionChecks, group.acName)
+		acs := admissioncheck.FindAdmissionCheck(group.local.Status.AdmissionChecks, group.acName)
 		if err := group.jobAdapter.SyncJob(ctx, w.client, group.remoteClients[reservingRemote].client, group.controllerKey, group.local.Name, w.origin); err != nil {
 			log.V(2).Error(err, "creating remote controller object", "remote", reservingRemote)
 			// We'll retry this in the next reconcile.
diff --git a/pkg/controller/admissionchecks/provisioning/controller.go b/pkg/controller/admissionchecks/provisioning/controller.go
index 4555477c6..aaa618672 100644
--- a/pkg/controller/admissionchecks/provisioning/controller.go
+++ b/pkg/controller/admissionchecks/provisioning/controller.go
@@ -236,7 +236,7 @@ func (c *Controller) syncOwnedProvisionRequest(
 		if !c.reqIsNeeded(wl, prc) {
 			continue
 		}
-		ac := workload.FindAdmissionCheck(wlInfo.checkStates, checkName)
+		ac := admissioncheck.FindAdmissionCheck(wlInfo.checkStates, checkName)
 		if ac != nil && ac.State == kueue.CheckStateReady {
 			log.V(2).Info("Skip syncing of the ProvReq for admission check which is Ready", "workload", klog.KObj(wl), "admissionCheck", checkName)
 			continue
@@ -593,7 +593,7 @@ func (c *Controller) syncCheckStates(
 			}
 		}
 
-		existingCondition := workload.FindAdmissionCheck(wlPatch.Status.AdmissionChecks, checkState.Name)
+		existingCondition := admissioncheck.FindAdmissionCheck(wlPatch.Status.AdmissionChecks, checkState.Name)
 		if existingCondition != nil && existingCondition.State != checkState.State {
 			message := fmt.Sprintf("Admission check %s updated state from %s to %s", checkState.Name, existingCondition.State, checkState.State)
 			if checkState.Message != "" {
diff --git a/pkg/scheduler/flavorassigner/flavorassigner.go b/pkg/scheduler/flavorassigner/flavorassigner.go
index c657c0b52..45c2b0cb1 100644
--- a/pkg/scheduler/flavorassigner/flavorassigner.go
+++ b/pkg/scheduler/flavorassigner/flavorassigner.go
@@ -541,7 +541,7 @@ func (a *FlavorAssigner) assignFlavors(log logr.Logger, counts []int32) Assignme
 				// No need to compute again.
 				continue
 			}
-			flavors, status := a.findFlavorForPodSetResource(log, psIDs, requests, resName, assignment.Usage.Quota)
+			flavors, status := a.findFlavorForPodSets(log, psIDs, requests, resName, assignment.Usage.Quota)
 			if status.IsError() || len(flavors) == 0 {
 				groupFlavors = nil
 				groupStatus = *status
@@ -618,12 +618,12 @@ func (a *Assignment) append(requests resources.Requests, psAssignment *PodSetAss
 	a.LastState.LastTriedFlavorIdx = append(a.LastState.LastTriedFlavorIdx, flavorIdx)
 }
 
-// findFlavorForPodSetResource finds the flavor which can satisfy the podSet request
+// findFlavorForPodSets finds the flavor which can satisfy all the PodSet requests
 // for all resources in the same group as resName.
 // Returns the chosen flavor, along with the information about resources that need to be borrowed.
 // If the flavor cannot be immediately assigned, it returns a status with
 // reasons or failure.
-func (a *FlavorAssigner) findFlavorForPodSetResource(
+func (a *FlavorAssigner) findFlavorForPodSets(
 	log logr.Logger,
 	psIDs []int,
 	requests resources.Requests,
@@ -657,43 +657,12 @@ func (a *FlavorAssigner) findFlavorForPodSetResource(
 	for ; idx < len(resourceGroup.Flavors); idx++ {
 		attemptedFlavorIdx = idx
 		fName := resourceGroup.Flavors[idx]
-		flavor, exist := a.resourceFlavors[fName]
-		if !exist {
-			log.Error(nil, "Flavor not found", "Flavor", fName)
-			status.appendf("flavor %s not found", fName)
-			continue
-		}
 
-		var flavorUnmatchMessage *string
-		for psIdx, psID := range psIDs {
-			if features.Enabled(features.TopologyAwareScheduling) {
-				ps := &a.wl.Obj.Spec.PodSets[psID]
-				flavorUnmatchMessage = checkPodSetAndFlavorMatchForTAS(a.cq, ps, flavor)
-				if flavorUnmatchMessage != nil {
-					log.Error(nil, *flavorUnmatchMessage)
-					break
-				}
-			}
-			podSpec := podSets[psIdx].Template.Spec
-			taint, untolerated := corev1helpers.FindMatchingUntoleratedTaint(flavor.Spec.NodeTaints, append(podSpec.Tolerations, flavor.Spec.Tolerations...), func(t *corev1.Taint) bool {
-				return t.Effect == corev1.TaintEffectNoSchedule || t.Effect == corev1.TaintEffectNoExecute
-			})
-			if untolerated {
-				flavorUnmatchMessage = ptr.To(fmt.Sprintf("untolerated taint %s in flavor %s", taint, fName))
-				break
+		if fit, err := a.checkFlavorForPodSets(log, fName, psIDs, podSets, selectors, status); !fit {
+			if err != nil {
+				status.err = err
+				return nil, status
 			}
-			selector := selectors[psIdx]
-			if match, err := selector.Match(&corev1.Node{ObjectMeta: metav1.ObjectMeta{Labels: flavor.Spec.NodeLabels}}); !match || err != nil {
-				if err != nil {
-					status.err = err
-					return nil, status
-				}
-				flavorUnmatchMessage = ptr.To(fmt.Sprintf("flavor %s doesn't match node affinity", fName))
-				break
-			}
-		}
-		if flavorUnmatchMessage != nil {
-			status.reasons = append(status.reasons, *flavorUnmatchMessage)
 			continue
 		}
 		assignments := make(ResourceAssignment, len(requests))
@@ -776,6 +745,51 @@ func (a *FlavorAssigner) findFlavorForPodSetResource(
 	return bestAssignment, status
 }
 
+func (a *FlavorAssigner) checkFlavorForPodSets(
+	log logr.Logger,
+	flavorName kueue.ResourceFlavorReference,
+	psIDs []int,
+	podSets []*kueue.PodSet,
+	selectors []nodeaffinity.RequiredNodeAffinity,
+	status *Status,
+) (bool, error) {
+	flavor, exist := a.resourceFlavors[flavorName]
+	if !exist {
+		log.Error(nil, "Flavor not found", "Flavor", flavorName)
+		status.appendf("flavor %s not found", flavorName)
+		return false, nil
+	}
+
+	for psIdx, psID := range psIDs {
+		if features.Enabled(features.TopologyAwareScheduling) {
+			ps := &a.wl.Obj.Spec.PodSets[psID]
+			if message := checkPodSetAndFlavorMatchForTAS(a.cq, ps, flavor); message != nil {
+				log.Error(nil, *message)
+				status.appendf("%s", *message)
+				return false, nil
+			}
+		}
+		podSpec := podSets[psIdx].Template.Spec
+		taint, untolerated := corev1helpers.FindMatchingUntoleratedTaint(flavor.Spec.NodeTaints, append(podSpec.Tolerations, flavor.Spec.Tolerations...), func(t *corev1.Taint) bool {
+			return t.Effect == corev1.TaintEffectNoSchedule || t.Effect == corev1.TaintEffectNoExecute
+		})
+		if untolerated {
+			status.appendf("untolerated taint %s in flavor %s", taint, flavorName)
+			return false, nil
+		}
+		selector := selectors[psIdx]
+		if match, err := selector.Match(&corev1.Node{ObjectMeta: metav1.ObjectMeta{Labels: flavor.Spec.NodeLabels}}); !match || err != nil {
+			if err != nil {
+				status.err = err
+				return false, err
+			}
+			status.appendf("flavor %s doesn't match node affinity", flavorName)
+			return false, nil
+		}
+	}
+	return true, nil
+}
+
 func shouldTryNextFlavor(representativeMode granularMode, flavorFungibility kueue.FlavorFungibility) bool {
 	policyPreempt := flavorFungibility.WhenCanPreempt
 	policyBorrow := flavorFungibility.WhenCanBorrow
diff --git a/pkg/util/admissioncheck/admissioncheck.go b/pkg/util/admissioncheck/admissioncheck.go
index b7db98faa..ea0d05fb4 100644
--- a/pkg/util/admissioncheck/admissioncheck.go
+++ b/pkg/util/admissioncheck/admissioncheck.go
@@ -172,3 +172,14 @@ func NewAdmissionChecks(cq *kueue.ClusterQueue) map[kueue.AdmissionCheckReferenc
 	}
 	return checks
 }
+
+// FindAdmissionCheck - returns a pointer to the check identified by checkName if found in checks.
+func FindAdmissionCheck(checks []kueue.AdmissionCheckState, checkName kueue.AdmissionCheckReference) *kueue.AdmissionCheckState {
+	for i := range checks {
+		if checks[i].Name == checkName {
+			return &checks[i]
+		}
+	}
+
+	return nil
+}
diff --git a/pkg/webhooks/workload_webhook.go b/pkg/webhooks/workload_webhook.go
index e72170077..af6184ca1 100644
--- a/pkg/webhooks/workload_webhook.go
+++ b/pkg/webhooks/workload_webhook.go
@@ -301,10 +301,8 @@ func validateAdmissionUpdate(new, old *kueue.Admission, path *field.Path) field.
 		if len(new.PodSetAssignments) != len(old.PodSetAssignments) {
 			return apivalidation.ValidateImmutableField(new, old, path)
 		}
-		// Allow to update (set) TopologyAssignment when DelayedTopologyRequest
-		// transitions from Pending to Ready.
+		// Allow to update (set) TopologyAssignments
 		for i := range new.PodSetAssignments {
-			// allow to update TopologyAssignments
 			old.PodSetAssignments[i].TopologyAssignment = new.PodSetAssignments[i].TopologyAssignment
 			old.PodSetAssignments[i].DelayedTopologyRequest = new.PodSetAssignments[i].DelayedTopologyRequest
 		}
diff --git a/pkg/workload/admissionchecks.go b/pkg/workload/admissionchecks.go
index 9e65dd087..44fadfda1 100644
--- a/pkg/workload/admissionchecks.go
+++ b/pkg/workload/admissionchecks.go
@@ -25,6 +25,7 @@ import (
 	"k8s.io/utils/clock"
 
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
+	"sigs.k8s.io/kueue/pkg/util/admissioncheck"
 )
 
 // SyncAdmittedCondition sync the state of the Admitted condition based on the
@@ -81,17 +82,6 @@ func SyncAdmittedCondition(w *kueue.Workload, now time.Time) bool {
 	return apimeta.SetStatusCondition(&w.Status.Conditions, newCondition)
 }
 
-// FindAdmissionCheck - returns a pointer to the check identified by checkName if found in checks.
-func FindAdmissionCheck(checks []kueue.AdmissionCheckState, checkName kueue.AdmissionCheckReference) *kueue.AdmissionCheckState {
-	for i := range checks {
-		if checks[i].Name == checkName {
-			return &checks[i]
-		}
-	}
-
-	return nil
-}
-
 // resetChecksOnEviction sets all AdmissionChecks to Pending
 func resetChecksOnEviction(w *kueue.Workload, now time.Time) {
 	checks := w.Status.AdmissionChecks
@@ -113,7 +103,7 @@ func SetAdmissionCheckState(checks *[]kueue.AdmissionCheckState, newCheck kueue.
 	if checks == nil {
 		return
 	}
-	existingCondition := FindAdmissionCheck(*checks, newCheck.Name)
+	existingCondition := admissioncheck.FindAdmissionCheck(*checks, newCheck.Name)
 	if existingCondition == nil {
 		if newCheck.LastTransitionTime.IsZero() {
 			newCheck.LastTransitionTime = metav1.NewTime(clock.Now())



Test patch:
diff --git a/pkg/scheduler/scheduler_test.go b/pkg/scheduler/scheduler_test.go
index 0ff12c4c6..703eb9136 100644
--- a/pkg/scheduler/scheduler_test.go
+++ b/pkg/scheduler/scheduler_test.go
@@ -3786,7 +3786,7 @@ func TestLastSchedulingContext(t *testing.T) {
 		cqs                            []kueue.ClusterQueue
 		admittedWorkloads              []kueue.Workload
 		workloads                      []kueue.Workload
-		deleteWorkloads                []kueue.Workload
+		deleteWorkloads                []client.ObjectKey
 		wantPreempted                  sets.Set[workload.Reference]
 		wantAdmissionsOnFirstSchedule  map[workload.Reference]kueue.Admission
 		wantAdmissionsOnSecondSchedule map[workload.Reference]kueue.Admission
@@ -3826,14 +3826,10 @@ func TestLastSchedulingContext(t *testing.T) {
 					Request(corev1.ResourceCPU, "20").
 					Obj(),
 			},
-			deleteWorkloads: []kueue.Workload{
-				*utiltesting.MakeWorkload("low-1", "default").
-					Queue("main").
-					Request(corev1.ResourceCPU, "50").
-					ReserveQuota(utiltesting.MakeAdmission("eng-alpha").Assignment(corev1.ResourceCPU, "on-demand", "50").Obj()).
-					Admitted(true).
-					Obj(),
-			},
+			deleteWorkloads: []client.ObjectKey{{
+				Namespace: metav1.NamespaceDefault,
+				Name:      "low-1",
+			}},
 			wantPreempted:                 sets.Set[workload.Reference]{},
 			wantAdmissionsOnFirstSchedule: map[workload.Reference]kueue.Admission{},
 			wantAdmissionsOnSecondSchedule: map[workload.Reference]kueue.Admission{
@@ -3860,8 +3856,7 @@ func TestLastSchedulingContext(t *testing.T) {
 					Request(corev1.ResourceCPU, "20").
 					Obj(),
 			},
-			deleteWorkloads: []kueue.Workload{},
-			wantPreempted:   sets.Set[workload.Reference]{},
+			wantPreempted: sets.Set[workload.Reference]{},
 			wantAdmissionsOnFirstSchedule: map[workload.Reference]kueue.Admission{
 				"default/workload1": *utiltesting.MakeAdmission("eng-cohort-beta").Assignment(corev1.ResourceCPU, "on-demand", "20").Obj(),
 				"default/borrower":  *utiltesting.MakeAdmission("eng-cohort-alpha").Assignment(corev1.ResourceCPU, "on-demand", "20").Obj(),
@@ -3893,8 +3888,7 @@ func TestLastSchedulingContext(t *testing.T) {
 					Request(corev1.ResourceCPU, "20").
 					Obj(),
 			},
-			deleteWorkloads: []kueue.Workload{},
-			wantPreempted:   sets.Set[workload.Reference]{},
+			wantPreempted: sets.Set[workload.Reference]{},
 			wantAdmissionsOnFirstSchedule: map[workload.Reference]kueue.Admission{
 				"default/workload": *utiltesting.MakeAdmission("eng-cohort-theta").Assignment(corev1.ResourceCPU, "spot", "20").Obj(),
 			},
@@ -3930,8 +3924,7 @@ func TestLastSchedulingContext(t *testing.T) {
 					Request(corev1.ResourceCPU, "20").
 					Obj(),
 			},
-			deleteWorkloads: []kueue.Workload{},
-			wantPreempted:   sets.Set[workload.Reference]{},
+			wantPreempted: sets.Set[workload.Reference]{},
 			wantAdmissionsOnFirstSchedule: map[workload.Reference]kueue.Admission{
 				"default/workload": *utiltesting.MakeAdmission("eng-cohort-theta").Assignment(corev1.ResourceCPU, "on-demand", "20").Obj(),
 			},
@@ -3964,7 +3957,10 @@ func TestLastSchedulingContext(t *testing.T) {
 					Request(corev1.ResourceCPU, "20").
 					Obj(),
 			},
-			deleteWorkloads:               []kueue.Workload{*utiltesting.MakeWorkload("placeholder-alpha", "default").Obj()},
+			deleteWorkloads: []client.ObjectKey{{
+				Namespace: metav1.NamespaceDefault,
+				Name:      "placeholder-alpha",
+			}},
 			wantPreempted:                 sets.New[workload.Reference]("default/placeholder-alpha"),
 			wantAdmissionsOnFirstSchedule: map[workload.Reference]kueue.Admission{},
 			wantAdmissionsOnSecondSchedule: map[workload.Reference]kueue.Admission{
@@ -4047,8 +4043,11 @@ func TestLastSchedulingContext(t *testing.T) {
 					Request(corev1.ResourceCPU, "22").
 					Obj(),
 			},
-			deleteWorkloads: []kueue.Workload{*utiltesting.MakeWorkload("alpha2", "default").Obj()},
-			wantPreempted:   sets.New[workload.Reference]("default/alpha2"),
+			deleteWorkloads: []client.ObjectKey{{
+				Namespace: metav1.NamespaceDefault,
+				Name:      "alpha2",
+			}},
+			wantPreempted: sets.New[workload.Reference]("default/alpha2"),
 			wantAdmissionsOnSecondSchedule: map[workload.Reference]kueue.Admission{
 				"default/alpha1": *utiltesting.MakeAdmission("eng-cohort-alpha").Assignment(corev1.ResourceCPU, "on-demand", "22").Obj(),
 				"default/alpha3": *utiltesting.MakeAdmission("eng-cohort-alpha").Assignment(corev1.ResourceCPU, "spot", "22").Obj(),
@@ -4130,16 +4129,21 @@ func TestLastSchedulingContext(t *testing.T) {
 				t.Errorf("Unexpected scheduled workloads (-want,+got):\n%s", diff)
 			}
 
-			for _, wl := range tc.deleteWorkloads {
-				err := cl.Delete(ctx, &wl)
+			for _, workloadReference := range tc.deleteWorkloads {
+				var workload kueue.Workload
+				err := cl.Get(ctx, workloadReference, &workload)
+				if err != nil {
+					t.Errorf("Unable to get workload: %v", err)
+				}
+				err = cl.Delete(ctx, &workload)
 				if err != nil {
 					t.Errorf("Delete workload failed: %v", err)
 				}
-				err = cqCache.DeleteWorkload(log, &wl)
+				err = cqCache.DeleteWorkload(log, &workload)
 				if err != nil {
 					t.Errorf("Delete workload failed: %v", err)
 				}
-				qManager.QueueAssociatedInadmissibleWorkloadsAfter(ctx, &wl, nil)
+				qManager.QueueAssociatedInadmissibleWorkloadsAfter(ctx, &workload, nil)
 			}
 
 			scheduler.schedule(ctx)
diff --git a/pkg/webhooks/workload_webhook_test.go b/pkg/webhooks/workload_webhook_test.go
index 27a795ce9..b9cbf45ac 100644
--- a/pkg/webhooks/workload_webhook_test.go
+++ b/pkg/webhooks/workload_webhook_test.go
@@ -376,7 +376,7 @@ func TestValidateWorkloadUpdate(t *testing.T) {
 				State:              kueue.CheckStateReady,
 			}).Obj(),
 		},
-		"TopologyAssignment cannot be mutated": {
+		"TopologyAssignment can be mutated": {
 			enableTopologyAwareScheduling: true,
 			before: testingutil.MakeWorkload(testWorkloadName, testWorkloadNamespace).
 				PodSets(
@@ -420,60 +420,22 @@ func TestValidateWorkloadUpdate(t *testing.T) {
 						Obj(),
 				).
 				Obj(),
-			wantErr: field.ErrorList{
-				field.Invalid(field.NewPath("status", "admission"), nil, ""),
-			},
+			wantErr: nil,
 		},
-		"TopologyAssignment cannot be set without delayedTopologyRequest": {
+		"PodSets cannot be removed from admission": {
 			enableTopologyAwareScheduling: true,
 			before: testingutil.MakeWorkload(testWorkloadName, testWorkloadNamespace).
 				PodSets(
 					*testingutil.MakePodSet("ps1", 3).Obj(),
+					*testingutil.MakePodSet("ps2", 3).Obj(),
 				).
 				ReserveQuota(
 					testingutil.MakeAdmission("cluster-queue").
 						PodSets(kueue.PodSetAssignment{
 							Name: "ps1",
 						}).
-						Obj(),
-				).
-				Obj(),
-			after: testingutil.MakeWorkload(testWorkloadName, testWorkloadNamespace).
-				PodSets(
-					*testingutil.MakePodSet("ps1", 3).Obj(),
-				).
-				ReserveQuota(
-					testingutil.MakeAdmission("cluster-queue").
 						PodSets(kueue.PodSetAssignment{
-							Name: "ps1",
-							TopologyAssignment: &kueue.TopologyAssignment{
-								Levels: []string{"level"},
-								Domains: []kueue.TopologyDomainAssignment{
-									{
-										Values: []string{"abc"},
-										Count:  3,
-									},
-								},
-							},
-						}).
-						Obj(),
-				).
-				Obj(),
-			wantErr: field.ErrorList{
-				field.Invalid(field.NewPath("status", "admission"), nil, ""),
-			},
-		},
-		"TopologyAssignment can be set if delayedTopologyRequest is Pending": {
-			enableTopologyAwareScheduling: true,
-			before: testingutil.MakeWorkload(testWorkloadName, testWorkloadNamespace).
-				PodSets(
-					*testingutil.MakePodSet("ps1", 3).Obj(),
-				).
-				ReserveQuota(
-					testingutil.MakeAdmission("cluster-queue").
-						PodSets(kueue.PodSetAssignment{
-							Name:                   "ps1",
-							DelayedTopologyRequest: ptr.To(kueue.DelayedTopologyRequestStatePending),
+							Name: "ps2",
 						}).
 						Obj(),
 				).
@@ -481,20 +443,12 @@ func TestValidateWorkloadUpdate(t *testing.T) {
 			after: testingutil.MakeWorkload(testWorkloadName, testWorkloadNamespace).
 				PodSets(
 					*testingutil.MakePodSet("ps1", 3).Obj(),
+					*testingutil.MakePodSet("ps2", 3).Obj(),
 				).
 				ReserveQuota(
 					testingutil.MakeAdmission("cluster-queue").
 						PodSets(kueue.PodSetAssignment{
 							Name: "ps1",
-							TopologyAssignment: &kueue.TopologyAssignment{
-								Levels: []string{"level"},
-								Domains: []kueue.TopologyDomainAssignment{
-									{
-										Values: []string{"abc"},
-										Count:  3,
-									},
-								},
-							},
 						}).
 						Obj(),
 				).
@@ -503,49 +457,22 @@ func TestValidateWorkloadUpdate(t *testing.T) {
 				field.Invalid(field.NewPath("status", "admission"), nil, ""),
 			},
 		},
-		"TopologyAssignment cannot be set if delayedTopologyRequest is Ready": {
+		"PodSets cannot be added to admission": {
 			enableTopologyAwareScheduling: true,
 			before: testingutil.MakeWorkload(testWorkloadName, testWorkloadNamespace).
 				PodSets(
 					*testingutil.MakePodSet("ps1", 3).Obj(),
-				).
-				ReserveQuota(
-					testingutil.MakeAdmission("cluster-queue").
-						PodSets(kueue.PodSetAssignment{
-							Name:                   "ps1",
-							DelayedTopologyRequest: ptr.To(kueue.DelayedTopologyRequestStateReady),
-						}).
-						Obj(),
-				).
-				Obj(),
-			after: testingutil.MakeWorkload(testWorkloadName, testWorkloadNamespace).
-				PodSets(
-					*testingutil.MakePodSet("ps1", 3).Obj(),
+					*testingutil.MakePodSet("ps2", 3).Obj(),
 				).
 				ReserveQuota(
 					testingutil.MakeAdmission("cluster-queue").
 						PodSets(kueue.PodSetAssignment{
 							Name: "ps1",
-							TopologyAssignment: &kueue.TopologyAssignment{
-								Levels: []string{"level"},
-								Domains: []kueue.TopologyDomainAssignment{
-									{
-										Values: []string{"abc"},
-										Count:  3,
-									},
-								},
-							},
 						}).
 						Obj(),
 				).
 				Obj(),
-			wantErr: field.ErrorList{
-				field.Invalid(field.NewPath("status", "admission"), nil, ""),
-			},
-		},
-		"PodSets cannot be removed from admission": {
-			enableTopologyAwareScheduling: true,
-			before: testingutil.MakeWorkload(testWorkloadName, testWorkloadNamespace).
+			after: testingutil.MakeWorkload(testWorkloadName, testWorkloadNamespace).
 				PodSets(
 					*testingutil.MakePodSet("ps1", 3).Obj(),
 					*testingutil.MakePodSet("ps2", 3).Obj(),
@@ -561,19 +488,6 @@ func TestValidateWorkloadUpdate(t *testing.T) {
 						Obj(),
 				).
 				Obj(),
-			after: testingutil.MakeWorkload(testWorkloadName, testWorkloadNamespace).
-				PodSets(
-					*testingutil.MakePodSet("ps1", 3).Obj(),
-					*testingutil.MakePodSet("ps2", 3).Obj(),
-				).
-				ReserveQuota(
-					testingutil.MakeAdmission("cluster-queue").
-						PodSets(kueue.PodSetAssignment{
-							Name: "ps1",
-						}).
-						Obj(),
-				).
-				Obj(),
 			wantErr: field.ErrorList{
 				field.Invalid(field.NewPath("status", "admission"), nil, ""),
 			},
@@ -619,6 +533,7 @@ func TestValidateWorkloadUpdate(t *testing.T) {
 	for name, tc := range testCases {
 		t.Run(name, func(t *testing.T) {
 			features.SetFeatureGateDuringTest(t, features.ElasticJobsViaWorkloadSlices, tc.enableElasticJobsFeature)
+			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.enableTopologyAwareScheduling)
 			errList := ValidateWorkloadUpdate(tc.after, tc.before, configapi.MultiKueueDispatcherModeAllAtOnce)
 			if diff := cmp.Diff(tc.wantErr, errList, cmpopts.IgnoreFields(field.Error{}, "Detail", "BadValue")); diff != "" {
 				t.Errorf("ValidateWorkloadUpdate() mismatch (-want +got):\n%s", diff)
diff --git a/pkg/workload/admissionchecks_test.go b/pkg/workload/admissionchecks_test.go
index 8e137b078..7a0699a26 100644
--- a/pkg/workload/admissionchecks_test.go
+++ b/pkg/workload/admissionchecks_test.go
@@ -29,6 +29,7 @@ import (
 
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 	"sigs.k8s.io/kueue/pkg/features"
+	"sigs.k8s.io/kueue/pkg/util/admissioncheck"
 	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
 )
 
@@ -536,7 +537,7 @@ func TestSetCheckState(t *testing.T) {
 			if tc.state.LastTransitionTime.IsZero() {
 				opts = append(opts, cmpopts.IgnoreFields(kueue.AdmissionCheckState{}, "LastTransitionTime"), cmpopts.EquateApproxTime(time.Second))
 
-				if updatedCheck := FindAdmissionCheck(gotStates, tc.state.Name); updatedCheck == nil {
+				if updatedCheck := admissioncheck.FindAdmissionCheck(gotStates, tc.state.Name); updatedCheck == nil {
 					t.Error("Cannot find the updated check state")
 				} else {
 					if diff := cmp.Diff(metav1.NewTime(now), updatedCheck.LastTransitionTime, opts...); diff != "" {
diff --git a/test/e2e/multikueue/e2e_test.go b/test/e2e/multikueue/e2e_test.go
index 27ab2121a..b3e3f771a 100644
--- a/test/e2e/multikueue/e2e_test.go
+++ b/test/e2e/multikueue/e2e_test.go
@@ -50,6 +50,7 @@ import (
 	podconstants "sigs.k8s.io/kueue/pkg/controller/jobs/pod/constants"
 	workloadraycluster "sigs.k8s.io/kueue/pkg/controller/jobs/raycluster"
 	workloadrayjob "sigs.k8s.io/kueue/pkg/controller/jobs/rayjob"
+	"sigs.k8s.io/kueue/pkg/util/admissioncheck"
 	utilpod "sigs.k8s.io/kueue/pkg/util/pod"
 	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
 	testingaw "sigs.k8s.io/kueue/pkg/util/testingjobs/appwrapper"
@@ -61,7 +62,6 @@ import (
 	testingpytorchjob "sigs.k8s.io/kueue/pkg/util/testingjobs/pytorchjob"
 	testingraycluster "sigs.k8s.io/kueue/pkg/util/testingjobs/raycluster"
 	testingrayjob "sigs.k8s.io/kueue/pkg/util/testingjobs/rayjob"
-	"sigs.k8s.io/kueue/pkg/workload"
 	"sigs.k8s.io/kueue/test/util"
 )
 
@@ -214,7 +214,7 @@ var _ = ginkgo.Describe("MultiKueue", func() {
 			ginkgo.By("Waiting to be admitted in worker2", func() {
 				gomega.Eventually(func(g gomega.Gomega) {
 					g.Expect(k8sManagerClient.Get(ctx, wlLookupKey, createdLeaderWorkload)).To(gomega.Succeed())
-					g.Expect(workload.FindAdmissionCheck(createdLeaderWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAc.Name))).To(gomega.BeComparableTo(&kueue.AdmissionCheckState{
+					g.Expect(admissioncheck.FindAdmissionCheck(createdLeaderWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAc.Name))).To(gomega.BeComparableTo(&kueue.AdmissionCheckState{
 						Name:    kueue.AdmissionCheckReference(multiKueueAc.Name),
 						State:   kueue.CheckStatePending,
 						Message: `The workload got reservation on "worker2"`,
@@ -410,7 +410,7 @@ var _ = ginkgo.Describe("MultiKueue", func() {
 			ginkgo.By("Waiting to be admitted in worker2", func() {
 				gomega.Eventually(func(g gomega.Gomega) {
 					g.Expect(k8sManagerClient.Get(ctx, wlLookupKey, createdLeaderWorkload)).To(gomega.Succeed())
-					g.Expect(workload.FindAdmissionCheck(createdLeaderWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAc.Name))).To(gomega.BeComparableTo(&kueue.AdmissionCheckState{
+					g.Expect(admissioncheck.FindAdmissionCheck(createdLeaderWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAc.Name))).To(gomega.BeComparableTo(&kueue.AdmissionCheckState{
 						Name:    kueue.AdmissionCheckReference(multiKueueAc.Name),
 						State:   kueue.CheckStatePending,
 						Message: `The workload got reservation on "worker2"`,
@@ -463,7 +463,7 @@ var _ = ginkgo.Describe("MultiKueue", func() {
 			ginkgo.By("Waiting to be admitted in worker2, and the manager's job unsuspended", func() {
 				gomega.Eventually(func(g gomega.Gomega) {
 					g.Expect(k8sManagerClient.Get(ctx, wlLookupKey, createdLeaderWorkload)).To(gomega.Succeed())
-					g.Expect(workload.FindAdmissionCheck(createdLeaderWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAc.Name))).To(gomega.BeComparableTo(&kueue.AdmissionCheckState{
+					g.Expect(admissioncheck.FindAdmissionCheck(createdLeaderWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAc.Name))).To(gomega.BeComparableTo(&kueue.AdmissionCheckState{
 						Name:    kueue.AdmissionCheckReference(multiKueueAc.Name),
 						State:   kueue.CheckStateReady,
 						Message: `The workload got reservation on "worker2"`,
@@ -1028,7 +1028,7 @@ func waitForJobAdmitted(wlLookupKey types.NamespacedName, acName, workerName str
 			Reason:  "Admitted",
 			Message: "The workload is admitted",
 		}, util.IgnoreConditionTimestampsAndObservedGeneration))
-		g.Expect(workload.FindAdmissionCheck(createdWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(acName))).To(gomega.BeComparableTo(&kueue.AdmissionCheckState{
+		g.Expect(admissioncheck.FindAdmissionCheck(createdWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(acName))).To(gomega.BeComparableTo(&kueue.AdmissionCheckState{
 			Name:    kueue.AdmissionCheckReference(acName),
 			State:   kueue.CheckStateReady,
 			Message: fmt.Sprintf(`The workload got reservation on "%s"`, workerName),
@@ -1060,7 +1060,7 @@ func ensurePodWorkloadsRunning(deployment *appsv1.Deployment, managerNs corev1.N
 		gomega.Expect(k8sManagerClient.Get(ctx, wlLookupKey, createdLeaderWorkload)).To(gomega.Succeed())
 
 		// By checking the assigned cluster we can discern which client to use
-		admissionCheckMessage := workload.FindAdmissionCheck(createdLeaderWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAc.Name)).Message
+		admissionCheckMessage := admissioncheck.FindAdmissionCheck(createdLeaderWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAc.Name)).Message
 		workerCluster := kubernetesClients[GetMultiKueueClusterNameFromAdmissionCheckMessage(admissionCheckMessage)]
 
 		// Worker pods should be in "Running" phase
diff --git a/test/integration/multikueue/dispatcher_test.go b/test/integration/multikueue/dispatcher_test.go
index 97ca671ef..acc0b6f5d 100644
--- a/test/integration/multikueue/dispatcher_test.go
+++ b/test/integration/multikueue/dispatcher_test.go
@@ -35,6 +35,7 @@ import (
 	"sigs.k8s.io/kueue/pkg/constants"
 	workloadjob "sigs.k8s.io/kueue/pkg/controller/jobs/job"
 	"sigs.k8s.io/kueue/pkg/features"
+	"sigs.k8s.io/kueue/pkg/util/admissioncheck"
 	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
 	testingjob "sigs.k8s.io/kueue/pkg/util/testingjobs/job"
 	"sigs.k8s.io/kueue/pkg/workload"
@@ -679,7 +680,7 @@ var _ = ginkgo.Describe("MultiKueueDispatcherAllAtOnce", ginkgo.Ordered, ginkgo.
 						Message: "The workload is admitted",
 					}, util.IgnoreConditionTimestampsAndObservedGeneration),
 				))
-				state := workload.FindAdmissionCheck(createdWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAC.Name))
+				state := admissioncheck.FindAdmissionCheck(createdWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAC.Name))
 				g.Expect(state).NotTo(gomega.BeNil())
 				g.Expect(state.State).To(gomega.Equal(kueue.CheckStateReady))
 			}, util.Timeout, util.Interval).Should(gomega.Succeed())
@@ -727,7 +728,7 @@ var _ = ginkgo.Describe("MultiKueueDispatcherAllAtOnce", ginkgo.Ordered, ginkgo.
 						Message: "The workload is admitted",
 					}, util.IgnoreConditionTimestampsAndObservedGeneration),
 				))
-				state := workload.FindAdmissionCheck(createdWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAC.Name))
+				state := admissioncheck.FindAdmissionCheck(createdWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAC.Name))
 				g.Expect(state).NotTo(gomega.BeNil())
 				g.Expect(state.State).To(gomega.Equal(kueue.CheckStateReady))
 			}, util.Timeout, util.Interval).Should(gomega.Succeed())
diff --git a/test/integration/multikueue/enabled_integration_test.go b/test/integration/multikueue/enabled_integration_test.go
index 369490ed5..ddc148716 100644
--- a/test/integration/multikueue/enabled_integration_test.go
+++ b/test/integration/multikueue/enabled_integration_test.go
@@ -39,10 +39,10 @@ import (
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 	workloadjob "sigs.k8s.io/kueue/pkg/controller/jobs/job"
 	workloadmpijob "sigs.k8s.io/kueue/pkg/controller/jobs/mpijob"
+	"sigs.k8s.io/kueue/pkg/util/admissioncheck"
 	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
 	testingjob "sigs.k8s.io/kueue/pkg/util/testingjobs/job"
 	testingmpijob "sigs.k8s.io/kueue/pkg/util/testingjobs/mpijob"
-	"sigs.k8s.io/kueue/pkg/workload"
 	"sigs.k8s.io/kueue/test/util"
 )
 
@@ -172,7 +172,7 @@ var _ = ginkgo.Describe("MultiKueue when not all integrations are enabled", gink
 
 			gomega.Eventually(func(g gomega.Gomega) {
 				g.Expect(managerTestCluster.client.Get(managerTestCluster.ctx, wlLookupKey, createdWorkload)).To(gomega.Succeed())
-				acs := workload.FindAdmissionCheck(createdWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAC.Name))
+				acs := admissioncheck.FindAdmissionCheck(createdWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAC.Name))
 				g.Expect(acs).NotTo(gomega.BeNil())
 				g.Expect(acs.State).To(gomega.Equal(kueue.CheckStatePending))
 				g.Expect(acs.Message).To(gomega.Equal(`The workload got reservation on "worker1"`))
@@ -271,7 +271,7 @@ var _ = ginkgo.Describe("MultiKueue when not all integrations are enabled", gink
 		ginkgo.By("checking the workload creation was rejected in the management cluster", func() {
 			managerWl := &kueue.Workload{}
 			gomega.Expect(managerTestCluster.client.Get(managerTestCluster.ctx, wlLookupKey, managerWl)).To(gomega.Succeed())
-			acs := workload.FindAdmissionCheck(managerWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAC.Name))
+			acs := admissioncheck.FindAdmissionCheck(managerWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAC.Name))
 			gomega.Expect(acs).To(gomega.BeComparableTo(&kueue.AdmissionCheckState{
 				Name:    kueue.AdmissionCheckReference(multiKueueAC.Name),
 				State:   kueue.CheckStateRejected,
diff --git a/test/integration/multikueue/jobs_test.go b/test/integration/multikueue/jobs_test.go
index 0a7a26520..e34a111dd 100644
--- a/test/integration/multikueue/jobs_test.go
+++ b/test/integration/multikueue/jobs_test.go
@@ -53,6 +53,7 @@ import (
 	workloadraycluster "sigs.k8s.io/kueue/pkg/controller/jobs/raycluster"
 	workloadrayjob "sigs.k8s.io/kueue/pkg/controller/jobs/rayjob"
 	"sigs.k8s.io/kueue/pkg/features"
+	"sigs.k8s.io/kueue/pkg/util/admissioncheck"
 	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
 	testingaw "sigs.k8s.io/kueue/pkg/util/testingjobs/appwrapper"
 	testingjob "sigs.k8s.io/kueue/pkg/util/testingjobs/job"
@@ -65,7 +66,6 @@ import (
 	testingrayjob "sigs.k8s.io/kueue/pkg/util/testingjobs/rayjob"
 	testingtfjob "sigs.k8s.io/kueue/pkg/util/testingjobs/tfjob"
 	testingxgboostjob "sigs.k8s.io/kueue/pkg/util/testingjobs/xgboostjob"
-	"sigs.k8s.io/kueue/pkg/workload"
 	"sigs.k8s.io/kueue/test/integration/framework"
 	"sigs.k8s.io/kueue/test/util"
 )
@@ -243,7 +243,7 @@ var _ = ginkgo.Describe("MultiKueue", ginkgo.Ordered, ginkgo.ContinueOnFailure,
 
 			gomega.Eventually(func(g gomega.Gomega) {
 				g.Expect(managerTestCluster.client.Get(managerTestCluster.ctx, wlLookupKey, createdWorkload)).To(gomega.Succeed())
-				acs := workload.FindAdmissionCheck(createdWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAC.Name))
+				acs := admissioncheck.FindAdmissionCheck(createdWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAC.Name))
 				g.Expect(acs).NotTo(gomega.BeNil())
 				g.Expect(acs.State).To(gomega.Equal(kueue.CheckStatePending))
 				g.Expect(acs.Message).To(gomega.Equal(`The workload got reservation on "worker1"`))
@@ -347,7 +347,7 @@ var _ = ginkgo.Describe("MultiKueue", ginkgo.Ordered, ginkgo.ContinueOnFailure,
 
 			gomega.Eventually(func(g gomega.Gomega) {
 				g.Expect(managerTestCluster.client.Get(managerTestCluster.ctx, wlLookupKey, createdWorkload)).To(gomega.Succeed())
-				acs := workload.FindAdmissionCheck(createdWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAC.Name))
+				acs := admissioncheck.FindAdmissionCheck(createdWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAC.Name))
 				g.Expect(acs).NotTo(gomega.BeNil())
 				g.Expect(acs.State).To(gomega.Equal(kueue.CheckStateReady))
 				g.Expect(acs.Message).To(gomega.Equal(`The workload got reservation on "worker1"`))
@@ -1062,7 +1062,7 @@ var _ = ginkgo.Describe("MultiKueue", ginkgo.Ordered, ginkgo.ContinueOnFailure,
 
 			gomega.Eventually(func(g gomega.Gomega) {
 				g.Expect(managerTestCluster.client.Get(managerTestCluster.ctx, wlLookupKey, createdWorkload)).To(gomega.Succeed())
-				acs := workload.FindAdmissionCheck(createdWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAC.Name))
+				acs := admissioncheck.FindAdmissionCheck(createdWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAC.Name))
 				g.Expect(acs).NotTo(gomega.BeNil())
 				g.Expect(acs.State).To(gomega.Equal(kueue.CheckStatePending))
 				g.Expect(acs.Message).To(gomega.Equal(`The workload got reservation on "worker1"`))
@@ -1179,7 +1179,7 @@ var _ = ginkgo.Describe("MultiKueue", ginkgo.Ordered, ginkgo.ContinueOnFailure,
 
 			gomega.Eventually(func(g gomega.Gomega) {
 				g.Expect(managerTestCluster.client.Get(managerTestCluster.ctx, wlLookupKey, createdWorkload)).To(gomega.Succeed())
-				acs := workload.FindAdmissionCheck(createdWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAC.Name))
+				acs := admissioncheck.FindAdmissionCheck(createdWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAC.Name))
 				g.Expect(acs).NotTo(gomega.BeNil())
 				g.Expect(acs.State).To(gomega.Equal(kueue.CheckStatePending))
 				g.Expect(acs.Message).To(gomega.Equal(`The workload got reservation on "worker1"`))
@@ -1463,7 +1463,7 @@ var _ = ginkgo.Describe("MultiKueue", ginkgo.Ordered, ginkgo.ContinueOnFailure,
 			gomega.Eventually(func(g gomega.Gomega) {
 				createdWorkload := &kueue.Workload{}
 				g.Expect(managerTestCluster.client.Get(managerTestCluster.ctx, wlLookupKey, createdWorkload)).To(gomega.Succeed())
-				acs := workload.FindAdmissionCheck(createdWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAC.Name))
+				acs := admissioncheck.FindAdmissionCheck(createdWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAC.Name))
 				g.Expect(acs).To(gomega.BeComparableTo(&kueue.AdmissionCheckState{
 					Name:  kueue.AdmissionCheckReference(multiKueueAC.Name),
 					State: kueue.CheckStatePending,
@@ -1625,7 +1625,7 @@ func admitWorkloadAndCheckWorkerCopies(acName string, wlLookupKey types.Namespac
 
 		gomega.Eventually(func(g gomega.Gomega) {
 			g.Expect(managerTestCluster.client.Get(managerTestCluster.ctx, wlLookupKey, createdWorkload)).To(gomega.Succeed())
-			acs := workload.FindAdmissionCheck(createdWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(acName))
+			acs := admissioncheck.FindAdmissionCheck(createdWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(acName))
 			g.Expect(acs).NotTo(gomega.BeNil())
 			g.Expect(acs.State).To(gomega.Equal(kueue.CheckStateReady))
 			g.Expect(acs.Message).To(gomega.Equal(`The workload got reservation on "worker2"`))
diff --git a/test/integration/multikueue/no_gc_test.go b/test/integration/multikueue/no_gc_test.go
index 1df02ffe5..08695ab50 100644
--- a/test/integration/multikueue/no_gc_test.go
+++ b/test/integration/multikueue/no_gc_test.go
@@ -33,6 +33,7 @@ import (
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 	workloadjob "sigs.k8s.io/kueue/pkg/controller/jobs/job"
 	"sigs.k8s.io/kueue/pkg/features"
+	"sigs.k8s.io/kueue/pkg/util/admissioncheck"
 	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
 	testingjob "sigs.k8s.io/kueue/pkg/util/testingjobs/job"
 	"sigs.k8s.io/kueue/pkg/workload"
@@ -254,7 +255,7 @@ var _ = ginkgo.Describe("MultiKueue no GC", ginkgo.Ordered, ginkgo.ContinueOnFai
 		ginkgo.By("setting workload reservation in worker1, AC state is updated in manager and worker2 wl is removed", func() {
 			gomega.Eventually(func(g gomega.Gomega) {
 				g.Expect(managerTestCluster.client.Get(managerTestCluster.ctx, wlLookupKey, createdWorkload)).To(gomega.Succeed())
-				acs := workload.FindAdmissionCheck(createdWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAC.Name))
+				acs := admissioncheck.FindAdmissionCheck(createdWorkload.Status.AdmissionChecks, kueue.AdmissionCheckReference(multiKueueAC.Name))
 				g.Expect(acs).NotTo(gomega.BeNil())
 				g.Expect(acs.State).To(gomega.Equal(kueue.CheckStatePending))
 			}, util.Timeout, util.Interval).Should(gomega.Succeed())
diff --git a/test/integration/singlecluster/controller/admissionchecks/provisioning/provisioning_test.go b/test/integration/singlecluster/controller/admissionchecks/provisioning/provisioning_test.go
index ee9b3aa73..c4a33a490 100644
--- a/test/integration/singlecluster/controller/admissionchecks/provisioning/provisioning_test.go
+++ b/test/integration/singlecluster/controller/admissionchecks/provisioning/provisioning_test.go
@@ -33,6 +33,7 @@ import (
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 	"sigs.k8s.io/kueue/pkg/constants"
 	"sigs.k8s.io/kueue/pkg/controller/admissionchecks/provisioning"
+	"sigs.k8s.io/kueue/pkg/util/admissioncheck"
 	"sigs.k8s.io/kueue/pkg/util/testing"
 	"sigs.k8s.io/kueue/pkg/workload"
 	"sigs.k8s.io/kueue/test/integration/framework"
@@ -270,7 +271,7 @@ var _ = ginkgo.Describe("Provisioning", ginkgo.Ordered, ginkgo.ContinueOnFailure
 			ginkgo.By("Checking that the ETA is propagated to workload", func() {
 				gomega.Eventually(func(g gomega.Gomega) {
 					g.Expect(k8sClient.Get(ctx, wlKey, &updatedWl)).To(gomega.Succeed())
-					state := workload.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
+					state := admissioncheck.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
 					g.Expect(state).NotTo(gomega.BeNil())
 					g.Expect(state.State).To(gomega.Equal(kueue.CheckStatePending))
 					g.Expect(state.Message).To(gomega.Equal("Not provisioned, ETA: 2024-02-22T10:36:40Z."))
@@ -292,7 +293,7 @@ var _ = ginkgo.Describe("Provisioning", ginkgo.Ordered, ginkgo.ContinueOnFailure
 			ginkgo.By("Checking the admission check", func() {
 				gomega.Eventually(func(g gomega.Gomega) {
 					g.Expect(k8sClient.Get(ctx, wlKey, &updatedWl)).To(gomega.Succeed())
-					state := workload.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
+					state := admissioncheck.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(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{
@@ -697,7 +698,7 @@ var _ = ginkgo.Describe("Provisioning", ginkgo.Ordered, ginkgo.ContinueOnFailure
 			ginkgo.By("Checking the admission check state indicates an inactive check", func() {
 				gomega.Eventually(func(g gomega.Gomega) {
 					g.Expect(k8sClient.Get(ctx, wlKey, &updatedWl)).To(gomega.Succeed())
-					state := workload.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
+					state := admissioncheck.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
 					g.Expect(state).NotTo(gomega.BeNil())
 					g.Expect(state.Message).To(gomega.Equal(provisioning.CheckInactiveMessage))
 				}, util.Timeout, util.Interval).Should(gomega.Succeed())
@@ -733,7 +734,7 @@ var _ = ginkgo.Describe("Provisioning", ginkgo.Ordered, ginkgo.ContinueOnFailure
 			ginkgo.By("Checking the admission check is ready", func() {
 				gomega.Eventually(func(g gomega.Gomega) {
 					g.Expect(k8sClient.Get(ctx, wlKey, &updatedWl)).To(gomega.Succeed())
-					state := workload.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
+					state := admissioncheck.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(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{
@@ -777,7 +778,7 @@ var _ = ginkgo.Describe("Provisioning", ginkgo.Ordered, ginkgo.ContinueOnFailure
 			ginkgo.By("Checking the admission check remains ready", func() {
 				gomega.Eventually(func(g gomega.Gomega) {
 					g.Expect(k8sClient.Get(ctx, wlKey, &updatedWl)).To(gomega.Succeed())
-					state := workload.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
+					state := admissioncheck.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
 					g.Expect(state).NotTo(gomega.BeNil())
 					g.Expect(state.State).To(gomega.Equal(kueue.CheckStateReady))
 				}, util.Timeout, util.Interval).Should(gomega.Succeed())
@@ -901,7 +902,7 @@ var _ = ginkgo.Describe("Provisioning", ginkgo.Ordered, ginkgo.ContinueOnFailure
 			ginkgo.By("Checking the AdmissionChecks are reset to Pending and remain this way", func() {
 				gomega.Consistently(func(g gomega.Gomega) {
 					g.Expect(k8sClient.Get(ctx, wlKey, &updatedWl)).To(gomega.Succeed())
-					check := workload.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
+					check := admissioncheck.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
 					g.Expect(check).NotTo(gomega.BeNil())
 					g.Expect(check.State).To(gomega.Equal(kueue.CheckStatePending), fmt.Sprintf("status: %v", updatedWl.Status))
 				}, util.ConsistentDuration, util.ShortInterval).Should(gomega.Succeed())
@@ -945,7 +946,7 @@ var _ = ginkgo.Describe("Provisioning", ginkgo.Ordered, ginkgo.ContinueOnFailure
 					g.Expect(k8sClient.Get(ctx, wlKey, &updatedWl)).To(gomega.Succeed())
 					_, evicted := workload.IsEvictedByAdmissionCheck(&updatedWl)
 					g.Expect(evicted).To(gomega.BeTrue())
-					check := workload.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
+					check := admissioncheck.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
 					g.Expect(check).NotTo(gomega.BeNil())
 					g.Expect(check.State).To(gomega.Equal(kueue.CheckStatePending))
 				}, util.Timeout, util.Interval).Should(gomega.Succeed())
@@ -1004,7 +1005,7 @@ var _ = ginkgo.Describe("Provisioning", ginkgo.Ordered, ginkgo.ContinueOnFailure
 			ginkgo.By("Checking the admission check is ready", func() {
 				gomega.Eventually(func(g gomega.Gomega) {
 					g.Expect(k8sClient.Get(ctx, wlKey, &updatedWl)).To(gomega.Succeed())
-					state := workload.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
+					state := admissioncheck.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
 					g.Expect(state).NotTo(gomega.BeNil())
 					g.Expect(state.State).To(gomega.Equal(kueue.CheckStateReady))
 				}, util.Timeout, util.Interval).Should(gomega.Succeed())
@@ -1048,7 +1049,7 @@ var _ = ginkgo.Describe("Provisioning", ginkgo.Ordered, ginkgo.ContinueOnFailure
 					g.Expect(k8sClient.Get(ctx, wlKey, &updatedWl)).To(gomega.Succeed())
 					_, evicted := workload.IsEvictedByAdmissionCheck(&updatedWl)
 					g.Expect(evicted).To(gomega.BeTrue())
-					check := workload.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
+					check := admissioncheck.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
 					g.Expect(check).NotTo(gomega.BeNil())
 					g.Expect(check.State).To(gomega.Equal(kueue.CheckStatePending))
 				}, util.Timeout, util.Interval).Should(gomega.Succeed())
@@ -1172,7 +1173,7 @@ var _ = ginkgo.Describe("Provisioning", ginkgo.Ordered, ginkgo.ContinueOnFailure
 					g.Expect(k8sClient.Get(ctx, wlKey, &updatedWl)).To(gomega.Succeed())
 					_, evicted := workload.IsEvictedByAdmissionCheck(&updatedWl)
 					g.Expect(evicted).To(gomega.BeTrue())
-					check := workload.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
+					check := admissioncheck.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
 					g.Expect(check).NotTo(gomega.BeNil())
 					g.Expect(check.State).To(gomega.Equal(kueue.CheckStatePending))
 				}, util.Timeout, util.Interval).Should(gomega.Succeed())
@@ -1231,7 +1232,7 @@ var _ = ginkgo.Describe("Provisioning", ginkgo.Ordered, ginkgo.ContinueOnFailure
 			ginkgo.By("Checking the admission check is ready", func() {
 				gomega.Eventually(func(g gomega.Gomega) {
 					g.Expect(k8sClient.Get(ctx, wlKey, &updatedWl)).To(gomega.Succeed())
-					state := workload.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
+					state := admissioncheck.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
 					g.Expect(state).NotTo(gomega.BeNil())
 					g.Expect(state.State).To(gomega.Equal(kueue.CheckStateReady))
 				}, util.Timeout, util.Interval).Should(gomega.Succeed())
@@ -1342,7 +1343,7 @@ var _ = ginkgo.Describe("Provisioning", ginkgo.Ordered, ginkgo.ContinueOnFailure
 					g.Expect(k8sClient.Get(ctx, wlKey, &updatedWl)).To(gomega.Succeed())
 					_, evicted := workload.IsEvictedByAdmissionCheck(&updatedWl)
 					g.Expect(evicted).To(gomega.BeTrue())
-					check := workload.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
+					check := admissioncheck.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
 					g.Expect(check).NotTo(gomega.BeNil())
 					g.Expect(check.State).To(gomega.Equal(kueue.CheckStatePending))
 				}, util.Timeout, util.Interval).Should(gomega.Succeed())
@@ -1391,7 +1392,7 @@ var _ = ginkgo.Describe("Provisioning", ginkgo.Ordered, ginkgo.ContinueOnFailure
 					g.Expect(k8sClient.Get(ctx, wlKey, &updatedWl)).To(gomega.Succeed())
 					_, evicted := workload.IsEvictedByAdmissionCheck(&updatedWl)
 					g.Expect(evicted).To(gomega.BeTrue())
-					check := workload.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
+					check := admissioncheck.FindAdmissionCheck(updatedWl.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
 					g.Expect(check).NotTo(gomega.BeNil())
 					g.Expect(check.State).To(gomega.Equal(kueue.CheckStatePending))
 				}, util.Timeout, util.Interval).Should(gomega.Succeed())
diff --git a/test/integration/singlecluster/controller/core/workload_controller_test.go b/test/integration/singlecluster/controller/core/workload_controller_test.go
index ee1a9ddce..4ebbc9140 100644
--- a/test/integration/singlecluster/controller/core/workload_controller_test.go
+++ b/test/integration/singlecluster/controller/core/workload_controller_test.go
@@ -33,6 +33,7 @@ import (
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 	"sigs.k8s.io/kueue/pkg/constants"
 	"sigs.k8s.io/kueue/pkg/features"
+	"sigs.k8s.io/kueue/pkg/util/admissioncheck"
 	"sigs.k8s.io/kueue/pkg/util/slices"
 	"sigs.k8s.io/kueue/pkg/util/testing"
 	"sigs.k8s.io/kueue/pkg/workload"
@@ -236,7 +237,7 @@ var _ = ginkgo.Describe("Workload controller", ginkgo.Ordered, ginkgo.ContinueOn
 			})
 
 			// save check2 condition
-			oldCheck2Cond := workload.FindAdmissionCheck(createdWl.Status.AdmissionChecks, "check2")
+			oldCheck2Cond := admissioncheck.FindAdmissionCheck(createdWl.Status.AdmissionChecks, "check2")
 			gomega.Expect(oldCheck2Cond).NotTo(gomega.BeNil())
 
 			ginkgo.By("updating the queue checks, the changes should propagate to the workload", func() {
@@ -256,7 +257,7 @@ var _ = ginkgo.Describe("Workload controller", ginkgo.Ordered, ginkgo.ContinueOn
 					})).Should(gomega.ConsistOf("check2", "check3"))
 				}, util.Timeout, util.Interval).Should(gomega.Succeed())
 
-				check2Cond := workload.FindAdmissionCheck(createdWl.Status.AdmissionChecks, "check2")
+				check2Cond := admissioncheck.FindAdmissionCheck(createdWl.Status.AdmissionChecks, "check2")
 				gomega.Expect(check2Cond).To(gomega.Equal(oldCheck2Cond))
 			})
 		})
diff --git a/test/integration/singlecluster/controller/jobs/job/job_controller_test.go b/test/integration/singlecluster/controller/jobs/job/job_controller_test.go
index bcf154db5..653e32e3f 100644
--- a/test/integration/singlecluster/controller/jobs/job/job_controller_test.go
+++ b/test/integration/singlecluster/controller/jobs/job/job_controller_test.go
@@ -3220,7 +3220,6 @@ var _ = ginkgo.Describe("Job with elastic jobs via workload-slices support", gin
 					g.Expect(workload.IsFinished(&workloads.Items[i])).Should(gomega.BeTrue())
 					continue
 				}
-				testJobWorkload = &workloads.Items[i]
 				util.ExpectWorkloadsToBeAdmitted(ctx, k8sClient, testJobWorkload)
 			}
 		}, util.Timeout, util.Interval).Should(gomega.Succeed())
diff --git a/test/integration/singlecluster/tas/tas_test.go b/test/integration/singlecluster/tas/tas_test.go
index e9e89708d..1e3b082e1 100644
--- a/test/integration/singlecluster/tas/tas_test.go
+++ b/test/integration/singlecluster/tas/tas_test.go
@@ -39,6 +39,7 @@ import (
 	"sigs.k8s.io/kueue/pkg/controller/admissionchecks/provisioning"
 	"sigs.k8s.io/kueue/pkg/controller/tas"
 	"sigs.k8s.io/kueue/pkg/features"
+	"sigs.k8s.io/kueue/pkg/util/admissioncheck"
 	utiltas "sigs.k8s.io/kueue/pkg/util/tas"
 	"sigs.k8s.io/kueue/pkg/util/testing"
 	testingnode "sigs.k8s.io/kueue/pkg/util/testingjobs/node"
@@ -2528,7 +2529,7 @@ var _ = ginkgo.Describe("Topology Aware Scheduling", ginkgo.Ordered, func() {
 				ginkgo.By("await for the check to be ready", func() {
 					gomega.Eventually(func(g gomega.Gomega) {
 						g.Expect(k8sClient.Get(ctx, wlKey, wl1)).To(gomega.Succeed())
-						state := workload.FindAdmissionCheck(wl1.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
+						state := admissioncheck.FindAdmissionCheck(wl1.Status.AdmissionChecks, kueue.AdmissionCheckReference(ac.Name))
 						g.Expect(state).NotTo(gomega.BeNil())
 						g.Expect(state.State).To(gomega.Equal(kueue.CheckStateReady))
 					}, util.Timeout, time.Millisecond).Should(gomega.Succeed())
diff --git a/test/util/util.go b/test/util/util.go
index 70c922479..466b5afaf 100644
--- a/test/util/util.go
+++ b/test/util/util.go
@@ -68,6 +68,7 @@ import (
 	podconstants "sigs.k8s.io/kueue/pkg/controller/jobs/pod/constants"
 	"sigs.k8s.io/kueue/pkg/metrics"
 	"sigs.k8s.io/kueue/pkg/scheduler/preemption"
+	"sigs.k8s.io/kueue/pkg/util/admissioncheck"
 	utiltas "sigs.k8s.io/kueue/pkg/util/tas"
 	"sigs.k8s.io/kueue/pkg/util/testing"
 	"sigs.k8s.io/kueue/pkg/workload"
@@ -747,7 +748,7 @@ func SetWorkloadsAdmissionCheck(ctx context.Context, k8sClient client.Client, wl
 	gomega.EventuallyWithOffset(1, func(g gomega.Gomega) {
 		g.Expect(k8sClient.Get(ctx, client.ObjectKeyFromObject(wl), &updatedWorkload)).To(gomega.Succeed())
 		if expectExisting {
-			currentCheck := workload.FindAdmissionCheck(updatedWorkload.Status.AdmissionChecks, check)
+			currentCheck := admissioncheck.FindAdmissionCheck(updatedWorkload.Status.AdmissionChecks, check)
 			g.Expect(currentCheck).NotTo(gomega.BeNil(), "the check %s was not found in %s", check, workload.Key(wl))
 			currentCheck.State = state
 		} else {