From 45811c3c847e76a1ebb046458057637700d52f75 Mon Sep 17 00:00:00 2001 From: Jason Parraga Date: Sat, 22 Feb 2025 12:48:13 -0800 Subject: [PATCH] Add unit test coverage Signed-off-by: Jason Parraga --- .../pluginmachinery/flytek8s/pod_helper.go | 6 +- .../flytek8s/pod_helper_test.go | 371 +++++++++++++----- 2 files changed, 269 insertions(+), 108 deletions(-) diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper.go index b52be30ebd..5d6d8f355e 100644 --- a/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper.go +++ b/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper.go @@ -691,7 +691,7 @@ func MergeWithBasePodTemplate(ctx context.Context, tCtx pluginsCore.TaskExecutio // magic values this method will merge containers that have matching names. func MergeBasePodSpecOntoTemplate(templatePodSpec *v1.PodSpec, basePodSpec *v1.PodSpec, primaryContainerName string, primaryInitContainerName string) (*v1.PodSpec, error) { if templatePodSpec == nil || basePodSpec == nil { - return nil, errors.New("templatePodSpec and basePodSpec must not be nil") + return nil, errors.New("neither the templatePodSpec or the basePodSpec can be nil") } // extract primaryContainerTemplate. The base should always contain the primary container. @@ -713,7 +713,7 @@ func MergeBasePodSpecOntoTemplate(templatePodSpec *v1.PodSpec, basePodSpec *v1.P for i := 0; i < len(templatePodSpec.InitContainers); i++ { if templatePodSpec.InitContainers[i].Name == defaultInitContainerTemplateName { defaultInitContainerTemplate = &templatePodSpec.InitContainers[i] - } else if primaryInitContainerName != primaryInitContainerTemplateName && templatePodSpec.InitContainers[i].Name == primaryInitContainerTemplateName { + } else if templatePodSpec.InitContainers[i].Name == primaryInitContainerTemplateName { primaryInitContainerTemplate = &templatePodSpec.InitContainers[i] } } @@ -764,7 +764,7 @@ func MergeBasePodSpecOntoTemplate(templatePodSpec *v1.PodSpec, basePodSpec *v1.P // Merge in the base container if mergedContainer == nil { mergedContainer = container.DeepCopy() - } else if container.Name != primaryContainerTemplateName { + } else { err := mergo.Merge(mergedContainer, container, mergo.WithOverride, mergo.WithAppendSlice) if err != nil { return nil, err diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper_test.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper_test.go index 66a53d0dcb..a782444d07 100644 --- a/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper_test.go +++ b/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper_test.go @@ -3,6 +3,7 @@ package flytek8s import ( "context" "encoding/json" + "errors" "fmt" "io/ioutil" "path/filepath" @@ -2087,134 +2088,294 @@ func TestMergeWithBasePodTemplate(t *testing.T) { }) } -func TestMergePodSpecs(t *testing.T) { - var priority int32 = 1 +func TestMergeBasePodSpecsOntoTemplate(t *testing.T) { - podSpec1, _ := MergeBasePodSpecOntoTemplate(nil, nil, "foo", "foo-init") - assert.Nil(t, podSpec1) + baseContainer1 := v1.Container{ + Name: "task-1", + Image: "task-image", + } - podSpec2, _ := MergeBasePodSpecOntoTemplate(nil, &v1.PodSpec{}, "foo", "foo-init") - assert.Nil(t, podSpec2) + baseContainer2 := v1.Container{ + Name: "task-2", + Image: "task-image", + } - podSpec := v1.PodSpec{ - Containers: []v1.Container{ - v1.Container{ - Name: "container", - VolumeMounts: []v1.VolumeMount{ + initContainer1 := v1.Container{ + Name: "task-init-1", + Image: "task-init-image", + } + + initContainer2 := v1.Container{ + Name: "task-init-2", + Image: "task-init-image", + } + + tests := []struct { + name string + templatePodSpec *v1.PodSpec + basePodSpec *v1.PodSpec + primaryContainerName string + primaryInitContainerName string + expectedResult *v1.PodSpec + expectedError error + }{ + { + name: "nil template", + templatePodSpec: nil, + basePodSpec: &v1.PodSpec{}, + expectedError: errors.New("neither the templatePodSpec or the basePodSpec can be nil"), + }, + { + name: "nil base", + templatePodSpec: &v1.PodSpec{}, + basePodSpec: nil, + expectedError: errors.New("neither the templatePodSpec or the basePodSpec can be nil"), + }, + { + name: "nil template and base", + templatePodSpec: nil, + basePodSpec: nil, + expectedError: errors.New("neither the templatePodSpec or the basePodSpec can be nil"), + }, + { + name: "template and base with no overlap", + templatePodSpec: &v1.PodSpec{ + SchedulerName: "templateScheduler", + }, + basePodSpec: &v1.PodSpec{ + ServiceAccountName: "baseServiceAccount", + }, + expectedResult: &v1.PodSpec{ + SchedulerName: "templateScheduler", + ServiceAccountName: "baseServiceAccount", + }, + }, + { + name: "template and base with overlap", + templatePodSpec: &v1.PodSpec{ + SchedulerName: "templateScheduler", + }, + basePodSpec: &v1.PodSpec{ + SchedulerName: "baseScheduler", + ServiceAccountName: "baseServiceAccount", + }, + expectedResult: &v1.PodSpec{ + SchedulerName: "baseScheduler", + ServiceAccountName: "baseServiceAccount", + }, + }, + { + name: "template with default containers and base with no containers", + templatePodSpec: &v1.PodSpec{ + Containers: []v1.Container{ + { + Name: "default", + Image: "default-image", + }, + }, + InitContainers: []v1.Container{ { - Name: "nccl", - MountPath: "abc", + Name: "default-init", + Image: "default-init-image", }, }, }, - v1.Container{ - Name: "bar", + basePodSpec: &v1.PodSpec{ + SchedulerName: "baseScheduler", + }, + expectedResult: &v1.PodSpec{ + SchedulerName: "baseScheduler", }, }, - InitContainers: []v1.Container{ - v1.Container{ - Name: "container-init", - VolumeMounts: []v1.VolumeMount{ + { + name: "template with no default containers and base containers", + templatePodSpec: &v1.PodSpec{}, + basePodSpec: &v1.PodSpec{ + Containers: []v1.Container{baseContainer1}, + InitContainers: []v1.Container{initContainer1}, + SchedulerName: "baseScheduler", + }, + expectedResult: &v1.PodSpec{ + Containers: []v1.Container{baseContainer1}, + InitContainers: []v1.Container{initContainer1}, + SchedulerName: "baseScheduler", + }, + }, + { + name: "template and base with matching containers", + templatePodSpec: &v1.PodSpec{ + Containers: []v1.Container{ + { + Name: "task-1", + Image: "default-task-image", + TerminationMessagePath: "/dev/template-termination-log", + }, + }, + InitContainers: []v1.Container{ { - Name: "nccl", - MountPath: "abc", + Name: "task-init-1", + Image: "default-task-init-image", + TerminationMessagePath: "/dev/template-init-termination-log", }, }, }, - v1.Container{ - Name: "bar-init", + basePodSpec: &v1.PodSpec{ + Containers: []v1.Container{baseContainer1}, + InitContainers: []v1.Container{initContainer1}, + SchedulerName: "baseScheduler", + }, + expectedResult: &v1.PodSpec{ + Containers: []v1.Container{ + { + Name: "task-1", + Image: "task-image", + TerminationMessagePath: "/dev/template-termination-log", + }, + }, + InitContainers: []v1.Container{ + { + Name: "task-init-1", + Image: "task-init-image", + TerminationMessagePath: "/dev/template-init-termination-log", + }, + }, + SchedulerName: "baseScheduler", }, }, - NodeSelector: map[string]string{ - "baz": "bar", - }, - Priority: &priority, - SchedulerName: "overrideScheduler", - Tolerations: []v1.Toleration{ - v1.Toleration{ - Key: "bar", + { + name: "template and base with no matching containers", + templatePodSpec: &v1.PodSpec{ + Containers: []v1.Container{ + { + Name: "not-matching", + Image: "default-task-image", + TerminationMessagePath: "/dev/template-termination-log", + }, + }, + InitContainers: []v1.Container{ + { + Name: "not-matching-init", + Image: "default-task-init-image", + TerminationMessagePath: "/dev/template-init-termination-log", + }, + }, }, - v1.Toleration{ - Key: "baz", + basePodSpec: &v1.PodSpec{ + Containers: []v1.Container{baseContainer1}, + InitContainers: []v1.Container{initContainer1}, + SchedulerName: "baseScheduler", + }, + expectedResult: &v1.PodSpec{ + Containers: []v1.Container{baseContainer1}, + InitContainers: []v1.Container{initContainer1}, + SchedulerName: "baseScheduler", }, }, - } - - defaultContainerTemplate := v1.Container{ - Name: defaultContainerTemplateName, - TerminationMessagePath: "/dev/default-termination-log", - } - - primaryContainerTemplate := v1.Container{ - Name: primaryContainerTemplateName, - TerminationMessagePath: "/dev/primary-termination-log", - } - - defaultInitContainerTemplate := v1.Container{ - Name: defaultInitContainerTemplateName, - TerminationMessagePath: "/dev/default-init-termination-log", - } - - primaryInitContainerTemplate := v1.Container{ - Name: primaryInitContainerTemplateName, - TerminationMessagePath: "/dev/primary-init-termination-log", - } - - podTemplateSpec := v1.PodSpec{ - Containers: []v1.Container{ - defaultContainerTemplate, - primaryContainerTemplate, - }, - InitContainers: []v1.Container{ - defaultInitContainerTemplate, - primaryInitContainerTemplate, - }, - HostNetwork: true, - NodeSelector: map[string]string{ - "foo": "bar", + { + name: "template with default containers and base with containers", + templatePodSpec: &v1.PodSpec{ + Containers: []v1.Container{ + { + Name: "default", + Image: "default-task-image", + TerminationMessagePath: "/dev/template-termination-log", + }, + }, + InitContainers: []v1.Container{ + { + Name: "default-init", + Image: "default-task-init-image", + TerminationMessagePath: "/dev/template-init-termination-log", + }, + }, + }, + basePodSpec: &v1.PodSpec{ + Containers: []v1.Container{baseContainer1, baseContainer2}, + InitContainers: []v1.Container{initContainer1, initContainer2}, + SchedulerName: "baseScheduler", + }, + expectedResult: &v1.PodSpec{ + Containers: []v1.Container{ + { + Name: "task-1", + Image: "task-image", + TerminationMessagePath: "/dev/template-termination-log", + }, + { + Name: "task-2", + Image: "task-image", + TerminationMessagePath: "/dev/template-termination-log", + }, + }, + InitContainers: []v1.Container{ + { + Name: "task-init-1", + Image: "task-init-image", + TerminationMessagePath: "/dev/template-init-termination-log", + }, + { + Name: "task-init-2", + Image: "task-init-image", + TerminationMessagePath: "/dev/template-init-termination-log", + }, + }, + SchedulerName: "baseScheduler", + }, }, - Tolerations: []v1.Toleration{ - v1.Toleration{ - Key: "foo", + { + name: "template with primary containers and base with containers", + templatePodSpec: &v1.PodSpec{ + Containers: []v1.Container{ + { + Name: "primary", + Image: "default-task-image", + TerminationMessagePath: "/dev/template-termination-log", + }, + }, + InitContainers: []v1.Container{ + { + Name: "primary-init", + Image: "default-task-init-image", + TerminationMessagePath: "/dev/template-init-termination-log", + }, + }, + }, + basePodSpec: &v1.PodSpec{ + Containers: []v1.Container{baseContainer1, baseContainer2}, + InitContainers: []v1.Container{initContainer1, initContainer2}, + SchedulerName: "baseScheduler", + }, + expectedResult: &v1.PodSpec{ + Containers: []v1.Container{ + { + Name: "task-1", + Image: "task-image", + TerminationMessagePath: "/dev/template-termination-log", + }, + baseContainer2, + }, + InitContainers: []v1.Container{ + { + Name: "task-init-1", + Image: "task-init-image", + TerminationMessagePath: "/dev/template-init-termination-log", + }, + initContainer2, + }, + SchedulerName: "baseScheduler", }, + primaryContainerName: "task-1", + primaryInitContainerName: "task-init-1", }, } - mergedPodSpec, err := MergeBasePodSpecOntoTemplate(&podTemplateSpec, &podSpec, podSpec.Containers[0].Name, podSpec.InitContainers[0].Name) - assert.Nil(t, err) - - // validate a PodTemplate-only field - assert.Equal(t, podTemplateSpec.HostNetwork, mergedPodSpec.HostNetwork) - // validate a PodSpec-only field - assert.Equal(t, podSpec.Priority, mergedPodSpec.Priority) - // validate an overwritten PodTemplate field - assert.Equal(t, podSpec.SchedulerName, mergedPodSpec.SchedulerName) - // validate a merged map - assert.Equal(t, len(podTemplateSpec.NodeSelector)+len(podSpec.NodeSelector), len(mergedPodSpec.NodeSelector)) - // validate an appended array - assert.Equal(t, len(podTemplateSpec.Tolerations)+len(podSpec.Tolerations), len(mergedPodSpec.Tolerations)) - - // validate first container - primaryContainer := mergedPodSpec.Containers[0] - assert.Equal(t, podSpec.Containers[0].Name, primaryContainer.Name) - assert.Equal(t, primaryContainerTemplate.TerminationMessagePath, primaryContainer.TerminationMessagePath) - assert.Equal(t, 1, len(primaryContainer.VolumeMounts)) - - // validate default container - defaultContainer := mergedPodSpec.Containers[1] - assert.Equal(t, podSpec.Containers[1].Name, defaultContainer.Name) - assert.Equal(t, defaultContainerTemplate.TerminationMessagePath, defaultContainer.TerminationMessagePath) - - // validate primary init container - primaryInitContainer := mergedPodSpec.InitContainers[0] - assert.Equal(t, podSpec.InitContainers[0].Name, primaryInitContainer.Name) - assert.Equal(t, primaryInitContainerTemplate.TerminationMessagePath, primaryInitContainer.TerminationMessagePath) - assert.Equal(t, 1, len(primaryInitContainer.VolumeMounts)) - - // validate default init container - defaultInitContainer := mergedPodSpec.InitContainers[1] - assert.Equal(t, podSpec.InitContainers[1].Name, defaultInitContainer.Name) - assert.Equal(t, defaultInitContainerTemplate.TerminationMessagePath, defaultInitContainer.TerminationMessagePath) + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result, mergeErr := MergeBasePodSpecOntoTemplate(tt.templatePodSpec, tt.basePodSpec, tt.primaryContainerName, tt.primaryInitContainerName) + assert.Equal(t, tt.expectedResult, result) + assert.Equal(t, tt.expectedError, mergeErr) + }) + } } func TestAddFlyteCustomizationsToContainer_SetConsoleUrl(t *testing.T) {