From 6ca43f371f346fff229fe45bf1502057b6ca2f68 Mon Sep 17 00:00:00 2001 From: Utku Ozdemir Date: Thu, 20 Nov 2025 15:37:11 +0100 Subject: [PATCH] test: pick UKI and non-UKI machines correctly Picker can retry if it couldn't pick enough machines, but we set the picked flags only once and never reset them. Change the approach to take it into account. Signed-off-by: Utku Ozdemir --- internal/integration/cluster_test.go | 33 +++++++++++++++++-------- internal/integration/kernelargs_test.go | 2 +- internal/integration/suites_test.go | 19 ++++++++------ 3 files changed, 35 insertions(+), 19 deletions(-) diff --git a/internal/integration/cluster_test.go b/internal/integration/cluster_test.go index 7e1c4850..a38c7ced 100644 --- a/internal/integration/cluster_test.go +++ b/internal/integration/cluster_test.go @@ -14,6 +14,7 @@ import ( "math/rand/v2" "os" "path/filepath" + "slices" "sync" "testing" "time" @@ -41,7 +42,7 @@ import ( // BeforeClusterCreateFunc is a function that is called before a cluster is created. type BeforeClusterCreateFunc func(ctx context.Context, t *testing.T, cli *client.Client, machineIDs []resource.ID) -type PickFilterFunc func(machineStatus *omni.MachineStatus) bool +type PickFilterFunc func(machineStatus *omni.MachineStatus, alreadyPicked []*omni.MachineStatus) bool // ClusterOptions are the options for cluster creation. // @@ -820,28 +821,40 @@ func pickUnallocatedMachines(ctx context.Context, t *testing.T, st state.State, result := make([]resource.ID, 0, count) if filterFunc == nil { - filterFunc = func(*omni.MachineStatus) bool { return true } + filterFunc = func(*omni.MachineStatus, []*omni.MachineStatus) bool { return true } } err := retry.Constant(time.Minute).RetryWithContext(ctx, func(ctx context.Context) error { + result = result[:0] // clear the result between retries + machineStatuses, err := safe.StateListAll[*omni.MachineStatus](ctx, st, state.WithLabelQuery(resource.LabelExists(omni.MachineStatusLabelAvailable))) if err != nil { return err } - machineIDs := make([]resource.ID, 0, machineStatuses.Len()) - for machineStatus := range machineStatuses.All() { - if filterFunc(machineStatus) { - machineIDs = append(machineIDs, machineStatus.Metadata().ID()) + shuffledMachineStatuses := slices.Collect(machineStatuses.All()) + rand.Shuffle(len(shuffledMachineStatuses), func(i, j int) { + shuffledMachineStatuses[i], shuffledMachineStatuses[j] = shuffledMachineStatuses[j], shuffledMachineStatuses[i] + }) + + pickedMachineStatuses := make([]*omni.MachineStatus, 0, count) + + for _, machineStatus := range shuffledMachineStatuses { + if filterFunc(machineStatus, pickedMachineStatuses) { + pickedMachineStatuses = append(pickedMachineStatuses, machineStatus) + } + + if len(pickedMachineStatuses) >= count { + break } } - if len(machineIDs) < count { - return retry.ExpectedErrorf("not enough machines: available %d, requested %d", len(machineIDs), count) + if len(pickedMachineStatuses) < count { + return retry.ExpectedErrorf("not enough machines: available %d, requested %d", len(pickedMachineStatuses), count) } - for _, j := range rand.Perm(len(machineIDs))[:count] { - result = append(result, machineIDs[j]) + for _, ms := range pickedMachineStatuses { + result = append(result, ms.Metadata().ID()) } return nil diff --git a/internal/integration/kernelargs_test.go b/internal/integration/kernelargs_test.go index 74257863..e89a1333 100644 --- a/internal/integration/kernelargs_test.go +++ b/internal/integration/kernelargs_test.go @@ -43,7 +43,7 @@ func testKernelArgsUpdate(t *testing.T, options *TestOptions) { SkipExtensionCheckOnCreate: options.SkipExtensionsCheckOnCreate, // Pick machines which are booted with UKI, as kernel args upgrades are only supported for them. - PickFilterFunc: func(ms *omni.MachineStatus) bool { + PickFilterFunc: func(ms *omni.MachineStatus, _ []*omni.MachineStatus) bool { return ms.TypedSpec().Value.GetSecurityState().GetBootedWithUki() }, }), diff --git a/internal/integration/suites_test.go b/internal/integration/suites_test.go index 0082b1f6..2b4d8255 100644 --- a/internal/integration/suites_test.go +++ b/internal/integration/suites_test.go @@ -1400,8 +1400,6 @@ Test Omni upgrades, the first half that runs on the previous Omni version omniClient := options.omniClient clusterName := "integration-omni-upgrades" - pickedUKI, pickedNonUKI := false, false - t.Run("ClusterShouldBeCreated", CreateCluster(t.Context(), omniClient, ClusterOptions{ Name: clusterName, ControlPlanes: 3, @@ -1420,22 +1418,27 @@ Test Omni upgrades, the first half that runs on the previous Omni version // Pick at least one UKI and one non-UKI machine, so that we cover both flows during the upgrade tests. // Schematic calculation differs between UKI and non-UKI machines, so it's important to test both. - PickFilterFunc: func(machineStatus *omni.MachineStatus) bool { + PickFilterFunc: func(machineStatus *omni.MachineStatus, alreadyPicked []*omni.MachineStatus) bool { isUKI := machineStatus.TypedSpec().Value.SecurityState.GetBootedWithUki() + pickedUKI, pickedNonUKI := false, false + + for _, ms := range alreadyPicked { + if ms.TypedSpec().Value.SecurityState.GetBootedWithUki() { + pickedUKI = true + } else { + pickedNonUKI = true + } + } if isUKI && !pickedUKI { // if it's UKI, and we haven't picked one yet, pick it - pickedUKI = true - return true } if !isUKI && !pickedNonUKI { // if it's non-UKI, and we haven't picked one yet, pick it - pickedNonUKI = true - return true } - return pickedUKI && pickedNonUKI // if we have already picked both types, allow any + return pickedUKI && pickedNonUKI // otherwise, pick only if we have both types already }, }))