diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/container_helper_test.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/container_helper_test.go index dcd356716d..6aab57f7b3 100644 --- a/flyteplugins/go/tasks/pluginmachinery/flytek8s/container_helper_test.go +++ b/flyteplugins/go/tasks/pluginmachinery/flytek8s/container_helper_test.go @@ -448,7 +448,7 @@ func TestToK8sContainer(t *testing.T) { assert.False(t, *container.SecurityContext.AllowPrivilegeEscalation) } -func getTemplateParametersForTest(resourceRequirements, platformResources *v1.ResourceRequirements, consoleURL string) template.Parameters { +func getTemplateParametersForTest(resourceRequirements, platformResources *v1.ResourceRequirements, includeConsoleURL bool, consoleURL string) template.Parameters { mockTaskExecMetadata := mocks.TaskExecutionMetadata{} mockTaskExecutionID := mocks.TaskExecutionID{} mockTaskExecutionID.OnGetUniqueNodeID().Return("unique_node_id") @@ -495,9 +495,10 @@ func getTemplateParametersForTest(resourceRequirements, platformResources *v1.Re mockOutputPath.OnGetPreviousCheckpointsPrefix().Return("/prev") return template.Parameters{ - TaskExecMetadata: &mockTaskExecMetadata, - Inputs: &mockInputReader, - OutputPath: &mockOutputPath, + TaskExecMetadata: &mockTaskExecMetadata, + Inputs: &mockInputReader, + OutputPath: &mockOutputPath, + IncludeConsoleURL: includeConsoleURL, } } @@ -509,7 +510,7 @@ func TestAddFlyteCustomizationsToContainer(t *testing.T) { Limits: v1.ResourceList{ v1.ResourceEphemeralStorage: resource.MustParse("2048Mi"), }, - }, nil, "") + }, nil, false, "") container := &v1.Container{ Command: []string{ "{{ .Input }}", @@ -557,7 +558,7 @@ func TestAddFlyteCustomizationsToContainer_Resources(t *testing.T) { Limits: v1.ResourceList{ v1.ResourceMemory: resource.MustParse("20"), }, - }, "") + }, false, "") err := AddFlyteCustomizationsToContainer(context.TODO(), templateParameters, ResourceCustomizationModeMergeExistingResources, container) assert.NoError(t, err) assert.True(t, container.Resources.Requests.Cpu().Equal(resource.MustParse("1"))) @@ -580,7 +581,7 @@ func TestAddFlyteCustomizationsToContainer_Resources(t *testing.T) { Limits: v1.ResourceList{ v1.ResourceMemory: resource.MustParse("20"), }, - }, "") + }, false, "") err := AddFlyteCustomizationsToContainer(context.TODO(), templateParameters, ResourceCustomizationModeMergeExistingResources, container) assert.NoError(t, err) assert.True(t, container.Resources.Requests.Cpu().Equal(resource.MustParse("1"))) @@ -615,7 +616,7 @@ func TestAddFlyteCustomizationsToContainer_Resources(t *testing.T) { v1.ResourceCPU: resource.MustParse("10"), v1.ResourceMemory: resource.MustParse("20"), }, - }, "") + }, false, "") err := AddFlyteCustomizationsToContainer(context.TODO(), templateParameters, ResourceCustomizationModeMergeExistingResources, container) assert.NoError(t, err) assert.True(t, container.Resources.Requests.Cpu().Equal(resource.MustParse("10"))) @@ -652,7 +653,7 @@ func TestAddFlyteCustomizationsToContainer_Resources(t *testing.T) { templateParameters := getTemplateParametersForTest(&v1.ResourceRequirements{ Requests: overrideRequests, Limits: overrideLimits, - }, &v1.ResourceRequirements{}, "") + }, &v1.ResourceRequirements{}, false, "") err := AddFlyteCustomizationsToContainer(context.TODO(), templateParameters, ResourceCustomizationModeMergeExistingResources, container) assert.NoError(t, err) @@ -687,7 +688,7 @@ func TestAddFlyteCustomizationsToContainer_ValidateExistingResources(t *testing. v1.ResourceCPU: resource.MustParse("10"), v1.ResourceMemory: resource.MustParse("20"), }, - }, "") + }, false, "") err := AddFlyteCustomizationsToContainer(context.TODO(), templateParameters, ResourceCustomizationModeEnsureExistingResourcesInRange, container) assert.NoError(t, err) diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper_test.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper_test.go index 8a76070788..b4b7e2ca0d 100644 --- a/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper_test.go +++ b/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper_test.go @@ -2121,18 +2121,33 @@ func TestMergePodSpecs(t *testing.T) { func TestAddFlyteCustomizationsToContainer_SetConsoleUrl(t *testing.T) { tests := []struct { - name string - consoleURL string - expectedEnvVar *v1.EnvVar + name string + includeConsoleURL bool + consoleURL string + expectedEnvVar *v1.EnvVar }{ { - name: "console url is not set", - consoleURL: "", - expectedEnvVar: nil, + name: "do not include console url and console url is not set", + includeConsoleURL: false, + consoleURL: "", + expectedEnvVar: nil, }, { - name: "console url is set", - consoleURL: "gopher://flyte:65535/console", + name: "include console url but console url is not set", + includeConsoleURL: false, + consoleURL: "", + expectedEnvVar: nil, + }, + { + name: "do not include console url but console url is set", + includeConsoleURL: false, + consoleURL: "gopher://flyte:65535/console", + expectedEnvVar: nil, + }, + { + name: "include console url and console url is set", + includeConsoleURL: true, + consoleURL: "gopher://flyte:65535/console", expectedEnvVar: &v1.EnvVar{ Name: "FLYTE_EXECUTION_URL", Value: "gopher://flyte:65535/console/projects/p2/domains/d2/executions/n2/nodeId/unique_node_id-1/nodes", @@ -2149,7 +2164,7 @@ func TestAddFlyteCustomizationsToContainer_SetConsoleUrl(t *testing.T) { "{{ .OutputPrefix }}", }, } - templateParameters := getTemplateParametersForTest(&v1.ResourceRequirements{}, &v1.ResourceRequirements{}, tt.consoleURL) + templateParameters := getTemplateParametersForTest(&v1.ResourceRequirements{}, &v1.ResourceRequirements{}, tt.includeConsoleURL, tt.consoleURL) err := AddFlyteCustomizationsToContainer(context.TODO(), templateParameters, ResourceCustomizationModeAssignResources, container) assert.NoError(t, err) if tt.expectedEnvVar == nil {