spark-operator/internal/controller/sparkapplication/web_ui_test.go

/*
Copyright 2017 Google LLC

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

    https://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 sparkapplication

import (
	"context"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"

	"github.com/kubeflow/spark-operator/v2/api/v1beta2"
	"github.com/kubeflow/spark-operator/v2/pkg/common"
)

func TestConfigWebUI(t *testing.T) {
	ctx := context.Background()
	appBase := &v1beta2.SparkApplication{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "test-app",
			Namespace: "default",
		},
		Status: v1beta2.SparkApplicationStatus{
			SparkApplicationID: "test-app-1",
		},
	}

	tests := []struct {
		name     string
		options  Options
		app      *v1beta2.SparkApplication
		wantConf map[string]string
	}{
		{
			name: "ui service disabled",
			options: Options{
				EnableUIService:  false,
				IngressURLFormat: "ingress.example.com/{{ $appName }}",
			},
			app:      appBase.DeepCopy(),
			wantConf: nil,
		},
		{
			name: "ingress format without path",
			options: Options{
				EnableUIService:  true,
				IngressURLFormat: "test-app.ingress.example.com",
			},
			app:      appBase.DeepCopy(),
			wantConf: nil,
		},
		{
			name: "ingress format with path",
			options: Options{
				EnableUIService:  true,
				IngressURLFormat: "ingress.example.com/{{ $appNamespace }}/{{ $appName }}",
			},
			app: appBase.DeepCopy(),
			wantConf: map[string]string{
				common.SparkUIProxyBase:        "/default/test-app",
				common.SparkUIProxyRedirectURI: "/",
			},
		},
		{
			name: "existing config preserved",
			options: Options{
				EnableUIService:  true,
				IngressURLFormat: "ingress.example.com/{{ $appNamespace }}/{{ $appName }}",
			},
			app: func() *v1beta2.SparkApplication {
				app := appBase.DeepCopy()
				app.Spec.SparkConf = map[string]string{
					common.SparkUIProxyBase:        "/custom",
					common.SparkUIProxyRedirectURI: "/keep",
					"spark.executor.instances":     "2",
				}
				return app
			}(),
			wantConf: map[string]string{
				common.SparkUIProxyBase:        "/custom",
				common.SparkUIProxyRedirectURI: "/keep",
				"spark.executor.instances":     "2",
			},
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			reconciler := &Reconciler{options: tc.options}
			err := reconciler.configWebUI(ctx, tc.app)
			require.NoError(t, err)

			if tc.wantConf == nil {
				assert.Nil(t, tc.app.Spec.SparkConf)
				return
			}

			assert.Equal(t, tc.wantConf, tc.app.Spec.SparkConf)
		})
	}
}

func TestGetWebUIServicePortName(t *testing.T) {
	appBase := &v1beta2.SparkApplication{Spec: v1beta2.SparkApplicationSpec{}}

	tests := []struct {
		name string
		app  *v1beta2.SparkApplication
		want string
	}{
		{
			name: "default port name when options nil",
			app:  appBase.DeepCopy(),
			want: common.DefaultSparkWebUIPortName,
		},
		{
			name: "default port name when ServicePortName nil",
			app: func() *v1beta2.SparkApplication {
				app := appBase.DeepCopy()
				app.Spec.SparkUIOptions = &v1beta2.SparkUIConfiguration{}
				return app
			}(),
			want: common.DefaultSparkWebUIPortName,
		},
		{
			name: "custom service port name",
			app: func() *v1beta2.SparkApplication {
				app := appBase.DeepCopy()
				name := "custom-port"
				app.Spec.SparkUIOptions = &v1beta2.SparkUIConfiguration{ServicePortName: &name}
				return app
			}(),
			want: "custom-port",
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			got := getWebUIServicePortName(tc.app)
			assert.Equal(t, tc.want, got)
		})
	}
}

func TestGetWebUIServicePort(t *testing.T) {
	appBase := &v1beta2.SparkApplication{Spec: v1beta2.SparkApplicationSpec{}}

	tests := []struct {
		name string
		app  *v1beta2.SparkApplication
		want int32
	}{
		{
			name: "port derived from spark conf when options nil",
			app: func() *v1beta2.SparkApplication {
				app := appBase.DeepCopy()
				app.Spec.SparkConf = map[string]string{
					common.SparkUIPortKey: "4041",
				}
				return app
			}(),
			want: 4041,
		},
		{
			name: "custom service port from options",
			app: func() *v1beta2.SparkApplication {
				app := appBase.DeepCopy()
				port := int32(18080)
				app.Spec.SparkUIOptions = &v1beta2.SparkUIConfiguration{ServicePort: &port}
				return app
			}(),
			want: 18080,
		},
		{
			name: "invalid conf falls back to default",
			app: func() *v1beta2.SparkApplication {
				app := appBase.DeepCopy()
				app.Spec.SparkConf = map[string]string{
					common.SparkUIPortKey: "not-a-port",
				}
				return app
			}(),
			want: common.DefaultSparkWebUIPort,
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			got, err := getWebUIServicePort(tc.app)
			require.NoError(t, err)
			assert.Equal(t, tc.want, got)
		})
	}
}

func TestGetWebUITargetPort(t *testing.T) {
	appBase := &v1beta2.SparkApplication{Spec: v1beta2.SparkApplicationSpec{}}

	tests := []struct {
		name string
		app  *v1beta2.SparkApplication
		want int32
	}{
		{
			name: "target port from spark conf",
			app: func() *v1beta2.SparkApplication {
				app := appBase.DeepCopy()
				app.Spec.SparkConf = map[string]string{
					common.SparkUIPortKey: "4045",
				}
				return app
			}(),
			want: 4045,
		},
		{
			name: "invalid value returns default",
			app: func() *v1beta2.SparkApplication {
				app := appBase.DeepCopy()
				app.Spec.SparkConf = map[string]string{
					common.SparkUIPortKey: "invalid",
				}
				return app
			}(),
			want: common.DefaultSparkWebUIPort,
		},
		{
			name: "missing value returns default",
			app:  appBase.DeepCopy(),
			want: common.DefaultSparkWebUIPort,
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			got, err := getWebUITargetPort(tc.app)
			require.NoError(t, err)
			assert.Equal(t, tc.want, got)
		})
	}
}

func TestCreateWebUIService(t *testing.T) {
	scheme := runtime.NewScheme()
	require.NoError(t, corev1.AddToScheme(scheme))
	require.NoError(t, v1beta2.AddToScheme(scheme))

	ctx := context.Background()

	tests := []struct {
		name         string
		app          *v1beta2.SparkApplication
		wantPort     int32
		wantTarget   int32
		wantPortName string
		wantType     corev1.ServiceType
		wantLabels   map[string]string
		wantAnnots   map[string]string
	}{
		{
			name: "defaults with spark conf port override",
			app: func() *v1beta2.SparkApplication {
				app := &v1beta2.SparkApplication{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "ui-default",
						Namespace: "default",
					},
					Spec: v1beta2.SparkApplicationSpec{
						SparkConf: map[string]string{
							common.SparkUIPortKey: "4041",
						},
					},
				}
				return app
			}(),
			wantPort:     4041,
			wantTarget:   4041,
			wantPortName: common.DefaultSparkWebUIPortName,
			wantType:     corev1.ServiceTypeClusterIP,
			wantLabels: map[string]string{
				common.LabelSparkAppName: "ui-default",
			},
			wantAnnots: nil,
		},
		{
			name: "custom service port and annotations",
			app: func() *v1beta2.SparkApplication {
				port := int32(80)
				portName := "http-spark"
				app := &v1beta2.SparkApplication{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "ui-custom",
						Namespace: "default",
					},
					Spec: v1beta2.SparkApplicationSpec{
						SparkConf: map[string]string{
							common.SparkUIPortKey: "4045",
						},
						SparkUIOptions: &v1beta2.SparkUIConfiguration{
							ServicePort:        &port,
							ServicePortName:    &portName,
							ServiceAnnotations: map[string]string{"key": "value"},
							ServiceLabels:      map[string]string{"custom": "label"},
						},
					},
				}
				return app
			}(),
			wantPort:     80,
			wantTarget:   4045,
			wantPortName: "http-spark",
			wantType:     corev1.ServiceTypeClusterIP,
			wantLabels: map[string]string{
				"custom": "label",
			},
			wantAnnots: map[string]string{"key": "value"},
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			client := fake.NewClientBuilder().WithScheme(scheme).Build()
			reconciler := &Reconciler{client: client}

			svc, err := reconciler.createWebUIService(ctx, tc.app)
			require.NoError(t, err)

			assert.Equal(t, tc.wantPort, svc.servicePort)
			assert.Equal(t, tc.wantTarget, svc.targetPort.IntVal)
			assert.Equal(t, tc.wantPortName, svc.servicePortName)
			assert.Equal(t, tc.wantType, svc.serviceType)

			created := &corev1.Service{}
			err = client.Get(ctx, types.NamespacedName{Name: svc.serviceName, Namespace: tc.app.Namespace}, created)
			require.NoError(t, err)

			if tc.wantLabels != nil {
				assert.Equal(t, tc.wantLabels, created.Labels)
			} else {
				assert.Equal(t, map[string]string{common.LabelSparkAppName: tc.app.Name}, created.Labels)
			}

			if tc.wantAnnots != nil {
				assert.Equal(t, tc.wantAnnots, created.Annotations)
			} else {
				assert.Nil(t, created.Annotations)
			}

			require.Len(t, created.Spec.Ports, 1)
			assert.Equal(t, tc.wantPort, created.Spec.Ports[0].Port)
			assert.Equal(t, tc.wantPortName, created.Spec.Ports[0].Name)
			assert.Equal(t, tc.wantTarget, created.Spec.Ports[0].TargetPort.IntVal)
			assert.Equal(t, tc.wantType, created.Spec.Type)
			assert.Equal(t, map[string]string{
				common.LabelSparkAppName: tc.app.Name,
				common.LabelSparkRole:    common.SparkRoleDriver,
			}, created.Spec.Selector)
		})
	}
}