Repository: kubernetes/kubernetes
File: pkg/scheduler/backend/workloadmanager/workloadmanager_test.go (Go)

Code before:
		{
			name:                 "binding an unscheduled pod",
			oldPod:               pod,
			newPod:               assignedPod,
			expectInAllPods:      true,
			expectInAssignedPods: true,
		},
		{
			name:                 "binding an assumed pod",
			assumePod:            true,
			oldPod:               pod,
			newPod:               assignedPod,
			expectInAllPods:      true,
			expectInAssignedPods: true,
		},
		{
			name:   "updating a non-workload pod is a no-op",
			oldPod: nonWorkloadPod,
			newPod: updatedNonWorkloadPod,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			_, ctx := ktesting.NewTestContext(t)
			ctx, cancel := context.WithCancel(ctx)
			defer cancel()

			manager := New(ctx)

			manager.AddPod(tt.oldPod)
			if tt.assumePod {
				info, err := manager.PodGroupInfo(tt.oldPod.Namespace, tt.oldPod.Spec.WorkloadRef)
				if err != nil {
					t.Fatalf("Unexpected error getting pod group info: %v", err)
				}
				info.AssumePod(tt.oldPod.UID)
			}

			manager.UpdatePod(tt.oldPod, tt.newPod)

			gotPodGroups := len(manager.podGroupInfos)
			if gotPodGroups == 0 {
				if tt.expectInAllPods {
					t.Fatalf("Expected pod group, but got none")
				}
				return
			}
			if !tt.expectInAllPods {
				t.Fatalf("Expected no pod groups, but got %v", gotPodGroups)
			}

Review comment (suggestion):
Is it needed? It looks like `ktesting.NewTestContext` already initializes a cancellable context and cancels it on test cleanup: 

https://github.com/kubernetes/kubernetes/blob/c2618d48c0623bec76c0709bd01ea87eb6e0cde3/test/utils/ktesting/ktesting.go#L48-L49

https://github.com/kubernetes/kubernetes/blob/c2618d48c0623bec76c0709bd01ea87eb6e0cde3/test/utils/ktesting/tcontext.go#L103-L106

https://github.com/kubernetes/kubernetes/blob/c2618d48c0623bec76c0709bd01ea87eb6e0cde3/test/utils/ktesting/tcontext.go#L160-L164

Code after:
			newPod:               assignedPod,
			expectInAllPods:      true,
			expectInAssignedPods: true,
		},
		{
			name:                 "binding an assumed pod",
			assumePod:            true,
			oldPod:               pod,
			newPod:               assignedPod,
			expectInAllPods:      true,
			expectInAssignedPods: true,
		},
		{
			name:   "updating a non-workload pod is a no-op",
			oldPod: nonWorkloadPod,
			newPod: updatedNonWorkloadPod,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			logger, _ := ktesting.NewTestContext(t)

			manager := New(logger)

			manager.AddPod(tt.oldPod)
			if tt.assumePod {
				state, err := manager.PodGroupState(tt.oldPod.Namespace, tt.oldPod.Spec.WorkloadRef)
				if err != nil {
					t.Fatalf("Unexpected error getting pod group state: %v", err)
				}
				state.AssumePod(tt.oldPod.UID)
			}

			manager.UpdatePod(tt.oldPod, tt.newPod)

			gotPodGroups := len(manager.podGroupStates)
			if gotPodGroups == 0 {
				if tt.expectInAllPods {
					t.Fatalf("Expected pod group, but got none")
				}
				return
			}
			if !tt.expectInAllPods {
				t.Fatalf("Expected no pod groups, but got %v", gotPodGroups)
			}
			state, err := manager.PodGroupState(tt.newPod.Namespace, tt.newPod.Spec.WorkloadRef)
			if err != nil {
				t.Fatalf("Unexpected error getting pod group state: %v", err)
			}
			if inAll := state.AllPods().Has(tt.newPod.UID); inAll != tt.expectInAllPods {