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

/*
Copyright 2018 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"
	"fmt"
	"testing"

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

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

// filteredCacheClient wraps a client.Client and overrides Get to simulate a
// label-filtered informer cache. ConfigMaps that lack the required label are
// invisible to Get (returns NotFound), just like a real ByObject label selector
// would behave at runtime.
type filteredCacheClient struct {
	client.Client
	requiredLabel string
	requiredValue string
}

func (f *filteredCacheClient) Get(ctx context.Context, key client.ObjectKey, obj client.Object, opts ...client.GetOption) error {
	if err := f.Client.Get(ctx, key, obj, opts...); err != nil {
		return err
	}
	if cm, ok := obj.(*corev1.ConfigMap); ok {
		if cm.Labels[f.requiredLabel] != f.requiredValue {
			return errors.NewNotFound(schema.GroupResource{Group: "", Resource: "configmaps"}, key.Name)
		}
	}
	return nil
}

func TestConfigPrometheusMonitoring(t *testing.T) {
	type testcase struct {
		app                   *v1beta2.SparkApplication
		metricsProperties     string
		metricsPropertiesFile string
		prometheusConfig      string
		port                  string
		driverJavaOptions     string
		executorJavaOptions   string
	}

	fakeClient := fake.NewFakeClient()

	testFn := func(test testcase, t *testing.T) {
		err := configPrometheusMonitoring(context.TODO(), test.app, fakeClient)
		assert.NoError(t, err, "failed to configure Prometheus monitoring")

		configMapName := util.GetPrometheusConfigMapName(test.app)
		configMap := &corev1.ConfigMap{
			ObjectMeta: metav1.ObjectMeta{Name: configMapName, Namespace: test.app.Namespace},
		}
		err = fakeClient.Get(context.TODO(), client.ObjectKeyFromObject(configMap), configMap)
		assert.NoError(t, err, "failed to get ConfigMap %s", configMapName)

		assert.Equal(t, "true", configMap.Labels[common.LabelCreatedBySparkOperator],
			"ConfigMap %s should have LabelCreatedBySparkOperator label", configMapName)

		if test.app.Spec.Monitoring.Prometheus != nil && test.app.Spec.Monitoring.Prometheus.ConfigFile == nil &&
			test.app.Spec.Monitoring.MetricsPropertiesFile == nil {
			assert.Len(t, configMap.Data, 2, "expected 2 data items")
		}

		if test.app.Spec.Monitoring.Prometheus != nil && test.app.Spec.Monitoring.Prometheus.ConfigFile != nil &&
			test.app.Spec.Monitoring.MetricsPropertiesFile == nil {
			assert.Len(t, configMap.Data, 1, "expected 1 data item")
		}

		if test.app.Spec.Monitoring.Prometheus != nil && test.app.Spec.Monitoring.Prometheus.ConfigFile == nil &&
			test.app.Spec.Monitoring.MetricsPropertiesFile != nil {
			assert.Len(t, configMap.Data, 1, "expected 1 data item")
		}

		if test.app.Spec.Monitoring.Prometheus != nil && test.app.Spec.Monitoring.MetricsPropertiesFile == nil {
			assert.Equal(t, test.metricsProperties, configMap.Data[common.MetricsPropertiesKey], "metrics.properties mismatch")
		}

		if test.app.Spec.Monitoring.Prometheus != nil && test.app.Spec.Monitoring.Prometheus.ConfigFile == nil {
			assert.Equal(t, test.prometheusConfig, configMap.Data[common.PrometheusConfigKey], "prometheus.yaml mismatch")
		}

		if test.app.Spec.Monitoring.Prometheus != nil && test.app.Spec.Monitoring.ExposeDriverMetrics {
			assert.Len(t, test.app.Spec.Driver.Annotations, 3, "expected 3 driver annotations")
			assert.Equal(t, test.port, test.app.Spec.Driver.Annotations[common.PrometheusPortAnnotation], "java agent port mismatch")
			assert.Equal(t, test.driverJavaOptions, *test.app.Spec.Driver.JavaOptions, "driver Java options mismatch")

		}

		if test.app.Spec.Monitoring.ExposeExecutorMetrics {
			assert.Len(t, test.app.Spec.Executor.Annotations, 3, "expected 3 executor annotations")
			assert.Equal(t, test.port, test.app.Spec.Executor.Annotations[common.PrometheusPortAnnotation], "java agent port mismatch")
			assert.Equal(t, test.executorJavaOptions, *test.app.Spec.Executor.JavaOptions, "executor Java options mismatch")
		}

		if test.app.Spec.Monitoring.MetricsPropertiesFile != nil {
			assert.Equal(t, test.metricsPropertiesFile, test.app.Spec.SparkConf["spark.metrics.conf"], "sparkConf mismatch")
		}
	}

	testcases := []testcase{
		{
			app: &v1beta2.SparkApplication{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "app1",
					Namespace: "default",
				},
				Spec: v1beta2.SparkApplicationSpec{
					Monitoring: &v1beta2.MonitoringSpec{
						ExposeDriverMetrics:   true,
						ExposeExecutorMetrics: true,
						Prometheus: &v1beta2.PrometheusSpec{
							JmxExporterJar: "/prometheus/exporter.jar",
						},
					},
				},
			},
			metricsProperties:   common.DefaultMetricsProperties,
			prometheusConfig:    common.DefaultPrometheusConfiguration,
			port:                fmt.Sprintf("%d", common.DefaultPrometheusJavaAgentPort),
			driverJavaOptions:   "-javaagent:/prometheus/exporter.jar=8090:/etc/metrics/conf/prometheus.yaml",
			executorJavaOptions: "-javaagent:/prometheus/exporter.jar=8090:/etc/metrics/conf/prometheus.yaml",
		},
		{
			app: &v1beta2.SparkApplication{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "app2",
					Namespace: "default",
				},
				Spec: v1beta2.SparkApplicationSpec{
					Driver: v1beta2.DriverSpec{
						JavaOptions: ptr.To("-XX:+PrintGCDetails -XX:+PrintGCTimeStamps"),
					},
					Executor: v1beta2.ExecutorSpec{
						JavaOptions: ptr.To("-XX:+PrintGCDetails -XX:+PrintGCTimeStamps"),
					},
					Monitoring: &v1beta2.MonitoringSpec{
						ExposeDriverMetrics:   true,
						ExposeExecutorMetrics: true,
						MetricsProperties:     ptr.To("testcase2dummy"),
						Prometheus: &v1beta2.PrometheusSpec{
							JmxExporterJar: "/prometheus/exporter.jar",
							Port:           ptr.To[int32](8091),
							Configuration:  ptr.To("testcase2dummy"),
						},
					},
				},
			},
			metricsProperties:   "testcase2dummy",
			prometheusConfig:    "testcase2dummy",
			port:                "8091",
			driverJavaOptions:   "-XX:+PrintGCDetails -XX:+PrintGCTimeStamps -javaagent:/prometheus/exporter.jar=8091:/etc/metrics/conf/prometheus.yaml",
			executorJavaOptions: "-XX:+PrintGCDetails -XX:+PrintGCTimeStamps -javaagent:/prometheus/exporter.jar=8091:/etc/metrics/conf/prometheus.yaml",
		},
		{
			app: &v1beta2.SparkApplication{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "app2",
					Namespace: "default",
				},
				Spec: v1beta2.SparkApplicationSpec{
					Driver: v1beta2.DriverSpec{
						JavaOptions: ptr.To("-XX:+PrintGCDetails -XX:+PrintGCTimeStamps"),
					},
					Executor: v1beta2.ExecutorSpec{
						JavaOptions: ptr.To("-XX:+PrintGCDetails -XX:+PrintGCTimeStamps"),
					},
					Monitoring: &v1beta2.MonitoringSpec{
						ExposeDriverMetrics:   true,
						ExposeExecutorMetrics: true,
						MetricsProperties:     ptr.To("testcase3dummy"),
						Prometheus: &v1beta2.PrometheusSpec{
							JmxExporterJar: "/prometheus/exporter.jar",
							Port:           ptr.To[int32](8091),
							ConfigFile:     ptr.To("testcase3dummy.yaml"),
						},
					},
				},
			},
			metricsProperties:   "testcase3dummy",
			port:                "8091",
			driverJavaOptions:   "-XX:+PrintGCDetails -XX:+PrintGCTimeStamps -javaagent:/prometheus/exporter.jar=8091:testcase3dummy.yaml",
			executorJavaOptions: "-XX:+PrintGCDetails -XX:+PrintGCTimeStamps -javaagent:/prometheus/exporter.jar=8091:testcase3dummy.yaml",
		},
		{
			app: &v1beta2.SparkApplication{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "app2",
					Namespace: "default",
				},
				Spec: v1beta2.SparkApplicationSpec{
					Driver: v1beta2.DriverSpec{
						JavaOptions: ptr.To("-XX:+PrintGCDetails -XX:+PrintGCTimeStamps"),
					},
					Executor: v1beta2.ExecutorSpec{
						JavaOptions: ptr.To("-XX:+PrintGCDetails -XX:+PrintGCTimeStamps"),
					},
					Monitoring: &v1beta2.MonitoringSpec{
						ExposeDriverMetrics:   true,
						ExposeExecutorMetrics: true,
						MetricsPropertiesFile: ptr.To("/testcase4dummy/metrics.properties"),
						Prometheus: &v1beta2.PrometheusSpec{
							JmxExporterJar: "/prometheus/exporter.jar",
							Port:           ptr.To[int32](8091),
							ConfigFile:     ptr.To("testcase4dummy.yaml"),
						},
					},
				},
			},
			metricsPropertiesFile: "/testcase4dummy/metrics.properties",
			port:                  "8091",
			driverJavaOptions:     "-XX:+PrintGCDetails -XX:+PrintGCTimeStamps -javaagent:/prometheus/exporter.jar=8091:testcase4dummy.yaml",
			executorJavaOptions:   "-XX:+PrintGCDetails -XX:+PrintGCTimeStamps -javaagent:/prometheus/exporter.jar=8091:testcase4dummy.yaml",
		},
		{
			app: &v1beta2.SparkApplication{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "app2",
					Namespace: "default",
				},
				Spec: v1beta2.SparkApplicationSpec{
					Driver: v1beta2.DriverSpec{
						JavaOptions: ptr.To("-XX:+PrintGCDetails -XX:+PrintGCTimeStamps"),
					},
					Executor: v1beta2.ExecutorSpec{
						JavaOptions: ptr.To("-XX:+PrintGCDetails -XX:+PrintGCTimeStamps"),
					},
					Monitoring: &v1beta2.MonitoringSpec{
						ExposeDriverMetrics:   true,
						ExposeExecutorMetrics: true,
						MetricsPropertiesFile: ptr.To("/testcase5dummy/metrics.properties"),
						Prometheus: &v1beta2.PrometheusSpec{
							JmxExporterJar: "/prometheus/exporter.jar",
							Port:           ptr.To[int32](8091),
						},
					},
				},
			},
			metricsPropertiesFile: "/testcase5dummy/metrics.properties",
			prometheusConfig:      common.DefaultPrometheusConfiguration,
			port:                  "8091",
			driverJavaOptions:     "-XX:+PrintGCDetails -XX:+PrintGCTimeStamps -javaagent:/prometheus/exporter.jar=8091:/etc/metrics/conf/prometheus.yaml",
			executorJavaOptions:   "-XX:+PrintGCDetails -XX:+PrintGCTimeStamps -javaagent:/prometheus/exporter.jar=8091:/etc/metrics/conf/prometheus.yaml",
		},
		{
			app: &v1beta2.SparkApplication{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "driver-only",
					Namespace: "default",
				},
				Spec: v1beta2.SparkApplicationSpec{
					Monitoring: &v1beta2.MonitoringSpec{
						ExposeDriverMetrics:   true,
						ExposeExecutorMetrics: false,
						Prometheus: &v1beta2.PrometheusSpec{
							JmxExporterJar: "/prometheus/exporter.jar",
						},
					},
					Driver: v1beta2.DriverSpec{JavaOptions: ptr.To("testdummy")},
				},
			},
			metricsProperties: common.DefaultMetricsProperties,
			prometheusConfig:  common.DefaultPrometheusConfiguration,
			port:              fmt.Sprintf("%d", common.DefaultPrometheusJavaAgentPort),
			driverJavaOptions: "testdummy -javaagent:/prometheus/exporter.jar=8090:/etc/metrics/conf/prometheus.yaml",
		},
		{
			app: &v1beta2.SparkApplication{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "executor-only",
					Namespace: "default",
				},
				Spec: v1beta2.SparkApplicationSpec{
					Monitoring: &v1beta2.MonitoringSpec{
						ExposeDriverMetrics:   false,
						ExposeExecutorMetrics: true,
						Prometheus: &v1beta2.PrometheusSpec{
							JmxExporterJar: "/prometheus/exporter.jar",
						},
					},
					Executor: v1beta2.ExecutorSpec{JavaOptions: ptr.To("testdummy")},
				},
			},
			metricsProperties:   common.DefaultMetricsProperties,
			prometheusConfig:    common.DefaultPrometheusConfiguration,
			port:                fmt.Sprintf("%d", common.DefaultPrometheusJavaAgentPort),
			executorJavaOptions: "testdummy -javaagent:/prometheus/exporter.jar=8090:/etc/metrics/conf/prometheus.yaml",
		},
		{
			app: &v1beta2.SparkApplication{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "custom-port-name",
					Namespace: "default",
				},
				Spec: v1beta2.SparkApplicationSpec{
					Monitoring: &v1beta2.MonitoringSpec{
						ExposeDriverMetrics:   true,
						ExposeExecutorMetrics: true,
						Prometheus: &v1beta2.PrometheusSpec{
							JmxExporterJar: "/prometheus/exporter.jar",
							Port:           ptr.To[int32](1000),
							PortName:       ptr.To("metrics-port"),
						},
					},
				},
			},
			metricsProperties:   common.DefaultMetricsProperties,
			prometheusConfig:    common.DefaultPrometheusConfiguration,
			port:                "1000",
			driverJavaOptions:   "-javaagent:/prometheus/exporter.jar=1000:/etc/metrics/conf/prometheus.yaml",
			executorJavaOptions: "-javaagent:/prometheus/exporter.jar=1000:/etc/metrics/conf/prometheus.yaml",
		},
		{
			app: &v1beta2.SparkApplication{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "app2",
					Namespace: "default",
				},
				Spec: v1beta2.SparkApplicationSpec{
					Driver: v1beta2.DriverSpec{
						JavaOptions: ptr.To("-XX:+PrintGCDetails -XX:+PrintGCTimeStamps"),
					},
					Executor: v1beta2.ExecutorSpec{
						JavaOptions: ptr.To("-XX:+PrintGCDetails -XX:+PrintGCTimeStamps"),
					},
					Monitoring: &v1beta2.MonitoringSpec{
						ExposeDriverMetrics:   false,
						ExposeExecutorMetrics: false,
						MetricsProperties:     ptr.To("testcase2dummy"),
					},
				},
			},
			metricsProperties:   "testcase2dummy",
			prometheusConfig:    "",
			port:                "8090",
			driverJavaOptions:   "-XX:+PrintGCDetails -XX:+PrintGCTimeStamps ",
			executorJavaOptions: "-XX:+PrintGCDetails -XX:+PrintGCTimeStamps ",
		},
	}

	for _, test := range testcases {
		testFn(test, t)
	}
}

func TestConfigPrometheusMonitoring_UpgradeExistingConfigMapWithoutLabel(t *testing.T) {
	app := &v1beta2.SparkApplication{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "upgrade-app",
			Namespace: "default",
			UID:       "test-uid",
		},
		Spec: v1beta2.SparkApplicationSpec{
			Monitoring: &v1beta2.MonitoringSpec{
				ExposeDriverMetrics:   true,
				ExposeExecutorMetrics: true,
				Prometheus: &v1beta2.PrometheusSpec{
					JmxExporterJar: "/prometheus/exporter.jar",
				},
			},
		},
	}

	configMapName := util.GetPrometheusConfigMapName(app)

	// Pre-create a ConfigMap WITHOUT the LabelCreatedBySparkOperator label,
	// simulating a ConfigMap created by a previous version of the operator.
	existingCM := &corev1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{
			Name:      configMapName,
			Namespace: app.Namespace,
			Labels: map[string]string{
				"some-other-label": "should-be-preserved",
			},
		},
		Data: map[string]string{
			"old-key": "old-value",
		},
	}

	underlying := fake.NewFakeClient(existingCM)

	// Wrap the fake client to simulate a label-filtered informer cache:
	// Get returns NotFound for ConfigMaps that lack LabelCreatedBySparkOperator,
	// forcing the code through the Create -> AlreadyExists -> Patch path.
	filtered := &filteredCacheClient{
		Client:        underlying,
		requiredLabel: common.LabelCreatedBySparkOperator,
		requiredValue: "true",
	}

	err := configPrometheusMonitoring(context.TODO(), app, filtered)
	assert.NoError(t, err, "configPrometheusMonitoring should handle pre-existing ConfigMap without label")

	// Read back from the underlying (unfiltered) client to verify the patch was applied.
	updatedCM := &corev1.ConfigMap{}
	err = underlying.Get(context.TODO(), client.ObjectKeyFromObject(existingCM), updatedCM)
	assert.NoError(t, err, "failed to get updated ConfigMap")

	assert.Equal(t, "true", updatedCM.Labels[common.LabelCreatedBySparkOperator],
		"Pre-existing ConfigMap should have LabelCreatedBySparkOperator added after upgrade")

	assert.Equal(t, "should-be-preserved", updatedCM.Labels["some-other-label"],
		"Merge patch should preserve pre-existing labels")

	assert.Contains(t, updatedCM.Data, common.MetricsPropertiesKey,
		"ConfigMap data should be updated with metrics properties")
}