Files
kubernetes/pkg/scheduler/backend/workloadmanager/workloadmanager_test.go
2025-11-06 10:47:29 +00:00

279 lines
8.6 KiB
Go

/*
Copyright 2025 The Kubernetes Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package workloadmanager
import (
"testing"
v1 "k8s.io/api/core/v1"
st "k8s.io/kubernetes/pkg/scheduler/testing"
)
func TestWorkloadManager_AddPod(t *testing.T) {
p1 := st.MakePod().Namespace("ns1").Name("p1").UID("p1").
WorkloadRef(&v1.WorkloadReference{Name: "w1", PodGroup: "pg1"}).Obj()
// Assigned
p2 := st.MakePod().Namespace("ns1").Name("p2").UID("p2").Node("node1").
WorkloadRef(&v1.WorkloadReference{Name: "w1", PodGroup: "pg1"}).Obj()
// Different ns
p3 := st.MakePod().Namespace("ns2").Name("p3").UID("p3").
WorkloadRef(&v1.WorkloadReference{Name: "w1", PodGroup: "pg1"}).Obj()
nonWorkloadPod := st.MakePod().Namespace("ns1").Name("non-workload").Obj()
tests := []struct {
name string
initPods []*v1.Pod
podToAdd *v1.Pod
expectedPodGroups int
expectInAllPods bool
expectInUnscheduledPods bool
expectInAssumedPods bool
expectInAssignedPods bool
}{
{
name: "adding an unscheduled pod",
podToAdd: p1,
expectedPodGroups: 1,
expectInAllPods: true,
expectInUnscheduledPods: true,
},
{
name: "adding an assigned pod",
podToAdd: p2,
expectedPodGroups: 1,
expectInAllPods: true,
expectInAssignedPods: true,
},
{
name: "adding pod with different namespace",
initPods: []*v1.Pod{p1},
podToAdd: p3,
expectedPodGroups: 2,
expectInAllPods: true,
expectInUnscheduledPods: true,
},
{
name: "adding a non-workload pod is a no-op",
podToAdd: nonWorkloadPod,
expectedPodGroups: 0,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
manager := New()
for _, p := range tt.initPods {
manager.AddPod(p)
}
manager.AddPod(tt.podToAdd)
gotPodGroups := len(manager.podGroupInfos)
if gotPodGroups != tt.expectedPodGroups {
t.Fatalf("Expected %v pod group(s), got %v", tt.expectedPodGroups, gotPodGroups)
}
if gotPodGroups == 0 {
return
}
info, err := manager.PodGroupInfo(tt.podToAdd.Namespace, tt.podToAdd.Spec.WorkloadRef)
if err != nil {
t.Fatalf("Unexpected error getting pod group info: %v", err)
}
if inAll := info.AllPods().Has(tt.podToAdd.UID); inAll != tt.expectInAllPods {
t.Errorf("Unexpected AllPods state, want: %v, got: %v", tt.expectInAllPods, inAll)
}
if inAssumed := info.AssumedPods().Has(tt.podToAdd.UID); inAssumed != tt.expectInAssumedPods {
t.Errorf("Unexpected AssumedPods state, want: %v, got: %v", tt.expectInAssumedPods, inAssumed)
}
if inAssigned := info.AssignedPods().Has(tt.podToAdd.UID); inAssigned != tt.expectInAssignedPods {
t.Errorf("Unexpected AssignedPods state, want: %v, got: %v", tt.expectInAssignedPods, inAssigned)
}
})
}
}
func TestWorkloadManager_UpdatePod(t *testing.T) {
pod := st.MakePod().Namespace("ns1").Name("p1").UID("p1").
WorkloadRef(&v1.WorkloadReference{Name: "w1", PodGroup: "pg1"}).Obj()
updatedPod := st.MakePod().Namespace("ns1").Name("p1").UID("p1").Labels(map[string]string{"foo": "bar"}).
WorkloadRef(&v1.WorkloadReference{Name: "w1", PodGroup: "pg1"}).Obj()
assignedPod := st.MakePod().Namespace("ns1").Name("p2").UID("p2").Node("node1").
WorkloadRef(&v1.WorkloadReference{Name: "w1", PodGroup: "pg1"}).Obj()
updatedAssignedPod := st.MakePod().Namespace("ns1").Name("p2").UID("p2").Node("node1").Labels(map[string]string{"foo": "bar"}).
WorkloadRef(&v1.WorkloadReference{Name: "w1", PodGroup: "pg1"}).Obj()
nonWorkloadPod := st.MakePod().Namespace("ns1").Name("non-workload").Obj()
updatedNonWorkloadPod := st.MakePod().Namespace("ns1").Name("non-workload").Labels(map[string]string{"foo": "bar"}).Obj()
tests := []struct {
name string
assumePod bool
oldPod *v1.Pod
newPod *v1.Pod
expectInAllPods bool
expectInUnscheduledPods bool
expectInAssumedPods bool
expectInAssignedPods bool
}{
{
name: "updating an unscheduled pod",
oldPod: pod,
newPod: updatedPod,
expectInAllPods: true,
expectInUnscheduledPods: true,
},
{
name: "updating an assumed pod",
assumePod: true,
oldPod: pod,
newPod: updatedPod,
expectInAllPods: true,
expectInAssumedPods: true,
},
{
name: "updating an assigned pod",
oldPod: assignedPod,
newPod: updatedAssignedPod,
expectInAllPods: true,
expectInAssignedPods: true,
},
{
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) {
manager := New()
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)
}
info, err := manager.PodGroupInfo(tt.newPod.Namespace, tt.newPod.Spec.WorkloadRef)
if err != nil {
t.Fatalf("Unexpected error getting pod group info: %v", err)
}
if inAll := info.AllPods().Has(tt.newPod.UID); inAll != tt.expectInAllPods {
t.Errorf("Unexpected AllPods state, want: %v, got: %v", tt.expectInAllPods, inAll)
}
if inAssumed := info.AssumedPods().Has(tt.newPod.UID); inAssumed != tt.expectInAssumedPods {
t.Errorf("Unexpected AssumedPods state, want: %v, got: %v", tt.expectInAssumedPods, inAssumed)
}
if inAssigned := info.AssignedPods().Has(tt.newPod.UID); inAssigned != tt.expectInAssignedPods {
t.Errorf("Unexpected AssignedPods state, want: %v, got: %v", tt.expectInAssignedPods, inAssigned)
}
})
}
}
func TestWorkloadManager_DeletePod(t *testing.T) {
p1 := st.MakePod().Namespace("ns1").Name("p1").UID("p1").
WorkloadRef(&v1.WorkloadReference{Name: "w1", PodGroup: "pg1"}).Obj()
p2 := st.MakePod().Namespace("ns1").Name("p2").UID("p2").
WorkloadRef(&v1.WorkloadReference{Name: "w1", PodGroup: "pg1"}).Obj()
tests := []struct {
name string
initPods []*v1.Pod
podToDelete *v1.Pod
expectedPodGroups int
}{
{
name: "deleting a pod from a group with multiple pods",
initPods: []*v1.Pod{p1, p2},
podToDelete: p1,
expectedPodGroups: 1,
},
{
name: "deleting the last pod cleans up the info",
initPods: []*v1.Pod{p1},
podToDelete: p1,
},
{
name: "deleting a non-existent pod is a no-op",
podToDelete: p1,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
manager := New()
for _, p := range tt.initPods {
manager.AddPod(p)
}
manager.DeletePod(tt.podToDelete)
gotPodGroups := len(manager.podGroupInfos)
if gotPodGroups != tt.expectedPodGroups {
t.Fatalf("Expected %v pod group(s), got %v", tt.expectedPodGroups, gotPodGroups)
}
if gotPodGroups == 0 {
return
}
info, err := manager.PodGroupInfo(tt.podToDelete.Namespace, tt.podToDelete.Spec.WorkloadRef)
if err != nil {
t.Fatalf("Unexpected error getting pod group info: %v", err)
}
if len(info.AllPods()) == 0 {
t.Errorf("Expected AllPods to be non-empty")
}
if info.AllPods().Has(p1.UID) {
t.Errorf("Expected pod to be deleted")
}
})
}
}