Fix bug and add unit tests

caraml-dev · Jan 7, 2025 · c5056b7 · c5056b7
1 parent 219b51b
commit c5056b7
Show file tree

Hide file tree

Showing 2 changed files with 86 additions and 1 deletion.
diff --git a/api/turing/cluster/spark.go b/api/turing/cluster/spark.go
@@ -304,5 +304,5 @@ func getEnvVars(request *CreateSparkRequest) []apicorev1.EnvVar {
 		envVars = append(envVars, apicorev1.EnvVar{Name: ev, Value: os.Getenv(ev)})
 	}
 
-	return append(defaultEnvVars, getEnvVarFromRequest(request)...)
+	return append(envVars, getEnvVarFromRequest(request)...)
 }
diff --git a/api/turing/cluster/spark_test.go b/api/turing/cluster/spark_test.go
@@ -1,6 +1,7 @@
 package cluster
 
 import (
+	"os"
 	"testing"
 
 	apisparkv1beta2 "github.com/GoogleCloudPlatform/spark-on-k8s-operator/pkg/apis/sparkoperator.k8s.io/v1beta2"
@@ -124,6 +125,90 @@ func TestGetCPURequestAndLimit(t *testing.T) {
 	}
 }
 
+func TestGetEnvVars(t *testing.T) {
+	request := &CreateSparkRequest{
+		JobName:               jobName,
+		JobLabels:             jobLabels,
+		JobImageRef:           jobImageRef,
+		JobApplicationPath:    jobApplicationPath,
+		JobArguments:          jobArguments,
+		JobConfigMount:        batch.JobConfigMount,
+		DriverCPURequest:      cpuValue,
+		DriverMemoryRequest:   memoryValue,
+		ExecutorCPURequest:    cpuValue,
+		ExecutorMemoryRequest: memoryValue,
+		ExecutorReplica:       executorReplica,
+		ServiceAccountName:    serviceAccountName,
+		SparkInfraConfig:      sparkInfraConfig,
+		EnvVars:               &envVars,
+	}
+	tests := map[string]struct {
+		sparkInfraConfigAPIServerEnvVars []string
+		apiServerEnvVars                 []apicorev1.EnvVar
+		expectedEnvVars                  []apicorev1.EnvVar
+	}{
+		"api server env vars specified": {
+			[]string{"TEST_ENV_VAR_1"},
+			[]apicorev1.EnvVar{
+				{
+					Name:  "TEST_ENV_VAR_1",
+					Value: "TEST_VALUE_1",
+				},
+			},
+			[]apicorev1.EnvVar{
+				{
+					Name:  envServiceAccountPathKey,
+					Value: envServiceAccountPath,
+				},
+				{
+					Name:  "TEST_ENV_VAR_1",
+					Value: "TEST_VALUE_1",
+				},
+				{
+					Name:  "foo",
+					Value: barString,
+				},
+			},
+		},
+		"no api server env vars specified": {
+			[]string{},
+			[]apicorev1.EnvVar{
+				{
+					Name:  "TEST_ENV_VAR_1",
+					Value: "TEST_VALUE_1",
+				},
+			},
+			[]apicorev1.EnvVar{
+				{
+					Name:  envServiceAccountPathKey,
+					Value: envServiceAccountPath,
+				},
+				{
+					Name:  "foo",
+					Value: barString,
+				},
+			},
+		},
+	}
+	for name, tt := range tests {
+		t.Run(name, func(t *testing.T) {
+			for _, ev := range tt.apiServerEnvVars {
+				err := os.Setenv(ev.Name, ev.Value)
+				assert.NoError(t, err)
+			}
+
+			request.SparkInfraConfig.APIServerEnvVars = tt.sparkInfraConfigAPIServerEnvVars
+			envVars := getEnvVars(request)
+			assert.Equal(t, tt.expectedEnvVars, envVars)
+
+			for _, ev := range tt.apiServerEnvVars {
+				err := os.Unsetenv(ev.Name)
+				assert.NoError(t, err)
+			}
+		})
+	}
+}
+
 var (
 	jobName                  = "jobname"
 	jobImageRef              = "gojek/nosuchimage"