Merge pull request #122524 from AxeZhan/testPQupdate

[Scheduler] refactor TestPriorityQueue_Update

pkg/scheduler/internal/queue/scheduling_queue_test.go

@@ -862,105 +862,136 @@ func TestPriorityQueue_Pop(t *testing.T) {
 }
 
 func TestPriorityQueue_Update(t *testing.T) {
-	objs := []runtime.Object{highPriorityPodInfo.Pod, unschedulablePodInfo.Pod, medPriorityPodInfo.Pod}
 	c := testingclock.NewFakeClock(time.Now())
-	logger, ctx := ktesting.NewTestContext(t)
+
-	ctx, cancel := context.WithCancel(ctx)
+	tests := []struct {
-	defer cancel()
+		name  string
-	q := NewTestQueueWithObjects(ctx, newDefaultQueueSort(), objs, WithClock(c))
+		wantQ string
-	// add highPriorityPodInfo to activeQ.
+		// wantAddedToNominated is whether a Pod from the test case should be registered as a nominated Pod in the nominator.
-	q.Update(logger, nil, highPriorityPodInfo.Pod)
+		wantAddedToNominated bool
-	if _, exists, _ := q.activeQ.Get(newQueuedPodInfoForLookup(highPriorityPodInfo.Pod)); !exists {
+		// prepareFunc is the function called to prepare pods in the queue before the test case calls Update().
-		t.Errorf("Expected %v to be added to activeQ.", highPriorityPodInfo.Pod.Name)
+		// This function returns three values;
-	}
+		// - oldPod/newPod: each test will call Update() with these oldPod and newPod.
-	if len(q.nominator.nominatedPods) != 0 {
+		prepareFunc func(t *testing.T, logger klog.Logger, q *PriorityQueue) (oldPod, newPod *v1.Pod)
-		t.Errorf("Expected nomindatePods to be empty: %v", q.nominator)
+	}{
-	}
+		{
-	// Update highPriorityPodInfo and add a nominatedNodeName to it.
+			name:  "add highPriorityPodInfo to activeQ",
-	q.Update(logger, highPriorityPodInfo.Pod, highPriNominatedPodInfo.Pod)
+			wantQ: activeQ,
-	if q.activeQ.Len() != 1 {
+			prepareFunc: func(t *testing.T, logger klog.Logger, q *PriorityQueue) (oldPod, newPod *v1.Pod) {
-		t.Error("Expected only one item in activeQ.")
+				return nil, highPriorityPodInfo.Pod
-	}
+			},
-	if len(q.nominator.nominatedPods) != 1 {
+		},
-		t.Errorf("Expected one item in nomindatePods map: %v", q.nominator)
+		{
-	}
+			name:  "Update pod that didn't exist in the queue",
-	// Updating an unschedulable pod which is not in any of the two queues, should
+			wantQ: activeQ,
-	// add the pod to activeQ.
+			prepareFunc: func(t *testing.T, logger klog.Logger, q *PriorityQueue) (oldPod, newPod *v1.Pod) {
-	q.Update(logger, unschedulablePodInfo.Pod, unschedulablePodInfo.Pod)
+				updatedPod := medPriorityPodInfo.Pod.DeepCopy()
-	if _, exists, _ := q.activeQ.Get(newQueuedPodInfoForLookup(unschedulablePodInfo.Pod)); !exists {
+				updatedPod.Annotations["foo"] = "test"
-		t.Errorf("Expected %v to be added to activeQ.", unschedulablePodInfo.Pod.Name)
+				return medPriorityPodInfo.Pod, updatedPod
-	}
+			},
-	// Updating a pod that is already in activeQ, should not change it.
+		},
-	q.Update(logger, unschedulablePodInfo.Pod, unschedulablePodInfo.Pod)
+		{
-	if len(q.unschedulablePods.podInfoMap) != 0 {
+			name:                 "Update highPriorityPodInfo and add a nominatedNodeName to it",
-		t.Error("Expected unschedulablePods to be empty.")
+			wantQ:                activeQ,
-	}
+			wantAddedToNominated: true,
-	if _, exists, _ := q.activeQ.Get(newQueuedPodInfoForLookup(unschedulablePodInfo.Pod)); !exists {
+			prepareFunc: func(t *testing.T, logger klog.Logger, q *PriorityQueue) (oldPod, newPod *v1.Pod) {
-		t.Errorf("Expected: %v to be added to activeQ.", unschedulablePodInfo.Pod.Name)
+				return highPriorityPodInfo.Pod, highPriNominatedPodInfo.Pod
-	}
+			},
-	if p, err := q.Pop(logger); err != nil || p.Pod != highPriNominatedPodInfo.Pod {
+		},
-		t.Errorf("Expected: %v after Pop, but got: %v", highPriorityPodInfo.Pod.Name, p.Pod.Name)
+		{
+			name:  "When updating a pod that is already in activeQ, the pod should remain in activeQ after Update()",
+			wantQ: activeQ,
+			prepareFunc: func(t *testing.T, logger klog.Logger, q *PriorityQueue) (oldPod, newPod *v1.Pod) {
+				err := q.Update(logger, nil, highPriorityPodInfo.Pod)
+				if err != nil {
+					t.Errorf("add pod %s error: %v", highPriorityPodInfo.Pod.Name, err)
+				}
+				return highPriorityPodInfo.Pod, highPriorityPodInfo.Pod
+			},
+		},
+		{
+			name:  "When updating a pod that is in backoff queue and is still backing off, it will be updated in backoff queue",
+			wantQ: backoffQ,
+			prepareFunc: func(t *testing.T, logger klog.Logger, q *PriorityQueue) (oldPod, newPod *v1.Pod) {
+				podInfo := q.newQueuedPodInfo(medPriorityPodInfo.Pod)
+				if err := q.podBackoffQ.Add(podInfo); err != nil {
+					t.Errorf("adding pod to backoff queue error: %v", err)
+				}
+				return podInfo.Pod, podInfo.Pod
+			},
+		},
+		{
+			name:  "when updating a pod which is in unschedulable queue and is backing off, it will be moved to backoff queue",
+			wantQ: backoffQ,
+			prepareFunc: func(t *testing.T, logger klog.Logger, q *PriorityQueue) (oldPod, newPod *v1.Pod) {
+				q.unschedulablePods.addOrUpdate(q.newQueuedPodInfo(medPriorityPodInfo.Pod, "plugin"))
+				updatedPod := medPriorityPodInfo.Pod.DeepCopy()
+				updatedPod.Annotations["foo"] = "test"
+				return medPriorityPodInfo.Pod, updatedPod
+			},
+		},
+		{
+			name:  "when updating a pod which is in unschedulable queue and is not backing off, it will be moved to active queue",
+			wantQ: activeQ,
+			prepareFunc: func(t *testing.T, logger klog.Logger, q *PriorityQueue) (oldPod, newPod *v1.Pod) {
+				q.unschedulablePods.addOrUpdate(q.newQueuedPodInfo(medPriorityPodInfo.Pod, "plugin"))
+				updatedPod := medPriorityPodInfo.Pod.DeepCopy()
+				updatedPod.Annotations["foo"] = "test1"
+				// Move clock by podInitialBackoffDuration, so that pods in the unschedulablePods would pass the backing off,
+				// and the pods will be moved into activeQ.
+				c.Step(q.podInitialBackoffDuration)
+				return medPriorityPodInfo.Pod, updatedPod
+			},
+		},
 	}
 
-	// Updating a pod that is in backoff queue and it is still backing off
+	for _, tt := range tests {
-	// pod will not be moved to active queue, and it will be updated in backoff queue
+		t.Run(tt.name, func(t *testing.T) {
-	podInfo := q.newQueuedPodInfo(medPriorityPodInfo.Pod)
+			logger, ctx := ktesting.NewTestContext(t)
-	if err := q.podBackoffQ.Add(podInfo); err != nil {
+			objs := []runtime.Object{highPriorityPodInfo.Pod, unschedulablePodInfo.Pod, medPriorityPodInfo.Pod}
-		t.Errorf("adding pod to backoff queue error: %v", err)
+			ctx, cancel := context.WithCancel(ctx)
-	}
+			defer cancel()
-	q.Update(logger, podInfo.Pod, podInfo.Pod)
+			q := NewTestQueueWithObjects(ctx, newDefaultQueueSort(), objs, WithClock(c))
-	rawPodInfo, err := q.podBackoffQ.Pop()
-	podGotFromBackoffQ := rawPodInfo.(*framework.QueuedPodInfo).Pod
-	if err != nil || podGotFromBackoffQ != medPriorityPodInfo.Pod {
-		t.Errorf("Expected: %v after Pop, but got: %v", medPriorityPodInfo.Pod.Name, podGotFromBackoffQ.Name)
-	}
 
-	// To simulate the pod is failed in scheduling in the real world, Pop() the pod from activeQ before testing AddUnschedulableIfNotPresent.
+			oldPod, newPod := tt.prepareFunc(t, logger, q)
-	q.activeQ.Add(podInfo)
-	if p, err := q.Pop(logger); err != nil || p.Pod != medPriorityPodInfo.Pod {
-		t.Errorf("Expected: %v after Pop, but got: %v", medPriorityPodInfo.Pod.Name, p.Pod.Name)
-	}
-	err = q.AddUnschedulableIfNotPresent(logger, q.newQueuedPodInfo(medPriorityPodInfo.Pod, "plugin"), q.SchedulingCycle())
-	if err != nil {
-		t.Fatalf("unexpected error from AddUnschedulableIfNotPresent: %v", err)
-	}
-	if len(q.unschedulablePods.podInfoMap) != 1 {
-		t.Error("Expected unschedulablePods to be 1.")
-	}
-	updatedPod := medPriorityPodInfo.Pod.DeepCopy()
-	updatedPod.Annotations["foo"] = "test"
-	q.Update(logger, medPriorityPodInfo.Pod, updatedPod)
-	rawPodInfo, err = q.podBackoffQ.Pop()
-	podGotFromBackoffQ = rawPodInfo.(*framework.QueuedPodInfo).Pod
-	if err != nil || podGotFromBackoffQ != updatedPod {
-		t.Errorf("Expected: %v after Pop, but got: %v", updatedPod.Name, podGotFromBackoffQ.Name)
-	}
 
-	// To simulate the pod is failed in scheduling in the real world, Pop() the pod from activeQ before testing AddUnschedulableIfNotPresent.
+			if err := q.Update(logger, oldPod, newPod); err != nil {
-	err = q.activeQ.Add(podInfo)
+				t.Fatalf("unexpected error from Update: %v", err)
-	if err != nil {
+			}
-		t.Fatalf("unexpected error from activeQ.Add: %v", err)
+
-	}
+			var pInfo *framework.QueuedPodInfo
-	if p, err := q.Pop(logger); err != nil || p.Pod != medPriorityPodInfo.Pod {
+
-		t.Errorf("Expected: %v after Pop, but got: %v", medPriorityPodInfo.Pod.Name, p.Pod.Name)
+			// validate expected queue
-	}
+			if obj, exists, _ := q.podBackoffQ.Get(newQueuedPodInfoForLookup(newPod)); exists {
-	// updating a pod which is in unschedulable queue, and it is not backing off,
+				if tt.wantQ != backoffQ {
-	// we will move it to active queue
+					t.Errorf("expected pod %s not to be queued to backoffQ, but it was", newPod.Name)
-	err = q.AddUnschedulableIfNotPresent(logger, q.newQueuedPodInfo(medPriorityPodInfo.Pod, "plugin"), q.SchedulingCycle())
+				}
-	if err != nil {
+				pInfo = obj.(*framework.QueuedPodInfo)
-		t.Fatalf("unexpected error from AddUnschedulableIfNotPresent: %v", err)
+			}
-	}
+
-	if len(q.unschedulablePods.podInfoMap) != 1 {
+			if obj, exists, _ := q.activeQ.Get(newQueuedPodInfoForLookup(newPod)); exists {
-		t.Error("Expected unschedulablePods to be 1.")
+				if tt.wantQ != activeQ {
-	}
+					t.Errorf("expected pod %s not to be queued to activeQ, but it was", newPod.Name)
-	updatedPod = medPriorityPodInfo.Pod.DeepCopy()
+				}
-	updatedPod.Annotations["foo"] = "test1"
+				pInfo = obj.(*framework.QueuedPodInfo)
-	// Move clock by podInitialBackoffDuration, so that pods in the unschedulablePods would pass the backing off,
+			}
-	// and the pods will be moved into activeQ.
+
-	c.Step(q.podInitialBackoffDuration)
+			if pInfoFromUnsched := q.unschedulablePods.get(newPod); pInfoFromUnsched != nil {
-	q.Update(logger, medPriorityPodInfo.Pod, updatedPod)
+				if tt.wantQ != unschedulablePods {
-	if p, err := q.Pop(logger); err != nil || p.Pod != updatedPod {
+					t.Errorf("expected pod %s to not be queued to unschedulablePods, but it was", newPod.Name)
-		t.Errorf("Expected: %v after Pop, but got: %v", updatedPod.Name, p.Pod.Name)
+				}
+				pInfo = pInfoFromUnsched
+			}
+
+			if diff := cmp.Diff(newPod, pInfo.PodInfo.Pod); diff != "" {
+				t.Errorf("Unexpected updated pod diff (-want, +got): %s", diff)
+			}
+
+			if tt.wantAddedToNominated && len(q.nominator.nominatedPods) != 1 {
+				t.Errorf("Expected one item in nomindatePods map: %v", q.nominator)
+			}
+
+		})
 	}
 }