Add test for preemption of multiple flavors

Change-Id: I6e46a40663a5f7e7a2cf93504051327f0b0f4766

pkg/scheduler/preemption/preemption_test.go

@@ -489,6 +489,62 @@ func TestPreemption(t *testing.T) {
 				},
 			}),
 		},
+		"each podset preempts a different flavor": {
+			admitted: []kueue.Workload{
+				*utiltesting.MakeWorkload("low-alpha", "").
+					Priority(-1).
+					Request(corev1.ResourceMemory, "2Gi").
+					Admit(utiltesting.MakeAdmission("standalone").Flavor(corev1.ResourceMemory, "alpha").Obj()).
+					Obj(),
+				*utiltesting.MakeWorkload("low-beta", "").
+					Priority(-1).
+					Request(corev1.ResourceMemory, "2Gi").
+					Admit(utiltesting.MakeAdmission("standalone").Flavor(corev1.ResourceMemory, "beta").Obj()).
+					Obj(),
+			},
+			incoming: utiltesting.MakeWorkload("in", "").
+				PodSets([]kueue.PodSet{
+					{
+						Name:  "launcher",
+						Count: 1,
+						Spec: utiltesting.PodSpecForRequest(map[corev1.ResourceName]string{
+							corev1.ResourceMemory: "2Gi",
+						}),
+					},
+					{
+						Name:  "workers",
+						Count: 2,
+						Spec: utiltesting.PodSpecForRequest(map[corev1.ResourceName]string{
+							corev1.ResourceMemory: "1Gi",
+						}),
+					},
+				}).
+				Obj(),
+			targetCQ: "standalone",
+			assignment: flavorassigner.Assignment{
+				PodSets: []flavorassigner.PodSetAssignment{
+					{
+						Name: "launcher",
+						Flavors: flavorassigner.ResourceAssignment{
+							corev1.ResourceMemory: {
+								Name: "alpha",
+								Mode: flavorassigner.Preempt,
+							},
+						},
+					},
+					{
+						Name: "workers",
+						Flavors: flavorassigner.ResourceAssignment{
+							corev1.ResourceMemory: {
+								Name: "beta",
+								Mode: flavorassigner.Preempt,
+							},
+						},
+					},
+				},
+			},
+			wantPreempted: sets.New("/low-alpha", "/low-beta"),
+		},
 	}
 	for name, tc := range cases {
 		t.Run(name, func(t *testing.T) {
@@ -585,6 +641,6 @@ func singlePodSetAssignment(assignments flavorassigner.ResourceAssignment) flavo
 		PodSets: []flavorassigner.PodSetAssignment{{
 			Name:    kueue.DefaultPodSetName,
 			Flavors: assignments,
-		},
-		}}
+		}},
+	}
 }