Fix flyteplugins tests

Signed-off-by: Eduardo Apolinario <[email protected]>
flyteorg · Jun 11, 2024 · d4e4e9f · d4e4e9f
1 parent 9079312
commit d4e4e9f
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 19 deletions.
diff --git 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
@@ -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 {