update unit tests for wireproto annotation

Author: fealebenpae

controllers/operator/mongodbsearch_controller_test.go

@@ -100,6 +100,19 @@ func buildExpectedMongotConfig(search *searchv1.MongoDBSearch, mdbc *mdbcv1.Mong
 	if search.Spec.LogLevel != "" {
 		logLevel = string(search.Spec.LogLevel)
 	}
+
+	var wireprotoServer *mongot.ConfigWireproto
+	if search.IsWireprotoForced() {
+		wireprotoServer = &mongot.ConfigWireproto{
+			Address: fmt.Sprintf("0.0.0.0:%d", search.GetMongotWireprotoPort()),
+			Authentication: &mongot.ConfigAuthentication{
+				Mode:    "keyfile",
+				KeyFile: searchcontroller.TempKeyfilePath,
+			},
+			TLS: &mongot.ConfigWireprotoTLS{Mode: mongot.ConfigTLSModeDisabled},
+		}
+	}
+
 	return mongot.Config{
 		SyncSource: mongot.ConfigSyncSource{
 			ReplicaSet: mongot.ConfigReplicaSet{
@@ -119,6 +132,7 @@ func buildExpectedMongotConfig(search *searchv1.MongoDBSearch, mdbc *mdbcv1.Mong
 				Address: fmt.Sprintf("0.0.0.0:%d", search.GetMongotGrpcPort()),
 				TLS:     &mongot.ConfigGrpcTLS{Mode: mongot.ConfigTLSModeDisabled},
 			},
+			Wireproto: wireprotoServer,
 		},
 		Metrics: mongot.ConfigMetrics{
 			Enabled: true,
@@ -164,35 +178,69 @@ func TestMongoDBSearchReconcile_MissingSource(t *testing.T) {
 
 func TestMongoDBSearchReconcile_Success(t *testing.T) {
 	ctx := context.Background()
-	search := newMongoDBSearch("search", mock.TestNamespace, "mdb")
-	search.Spec.LogLevel = "WARN"
-
-	mdbc := newMongoDBCommunity("mdb", mock.TestNamespace)
-	reconciler, c := newSearchReconciler(mdbc, search)
 
-	res, err := reconciler.Reconcile(
-		ctx,
-		reconcile.Request{NamespacedName: types.NamespacedName{Name: search.Name, Namespace: search.Namespace}},
-	)
-	expected, _ := workflow.OK().ReconcileResult()
-	assert.NoError(t, err)
-	assert.Equal(t, expected, res)
-
-	svc := &corev1.Service{}
-	err = c.Get(ctx, search.SearchServiceNamespacedName(), svc)
-	assert.NoError(t, err)
+	tests := []struct {
+		name          string
+		withWireproto bool
+	}{
+		{
+			name:          "grpc only (default)",
+			withWireproto: false,
+		},
+		{
+			name:          "grpc + wireproto via annotation",
+			withWireproto: true,
+		},
+	}
 
-	cm := &corev1.ConfigMap{}
-	err = c.Get(ctx, search.MongotConfigConfigMapNamespacedName(), cm)
-	assert.NoError(t, err)
-	expectedConfig := buildExpectedMongotConfig(search, mdbc)
-	configYaml, err := yaml.Marshal(expectedConfig)
-	assert.NoError(t, err)
-	assert.Equal(t, string(configYaml), cm.Data[searchcontroller.MongotConfigFilename])
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			search := newMongoDBSearch("search", mock.TestNamespace, "mdb")
+			search.Spec.LogLevel = "WARN"
+			if tc.withWireproto {
+				if search.Annotations == nil {
+					search.Annotations = map[string]string{}
+				}
+				search.Annotations[searchv1.ForceWireprotoTransportAnnotation] = "true"
+			}
 
-	sts := &appsv1.StatefulSet{}
-	err = c.Get(ctx, search.StatefulSetNamespacedName(), sts)
-	assert.NoError(t, err)
+			mdbc := newMongoDBCommunity("mdb", mock.TestNamespace)
+			reconciler, c := newSearchReconciler(mdbc, search)
+
+			res, err := reconciler.Reconcile(
+				ctx,
+				reconcile.Request{NamespacedName: types.NamespacedName{Name: search.Name, Namespace: search.Namespace}},
+			)
+			expected, _ := workflow.OK().ReconcileResult()
+			assert.NoError(t, err)
+			assert.Equal(t, expected, res)
+
+			svc := &corev1.Service{}
+			err = c.Get(ctx, search.SearchServiceNamespacedName(), svc)
+			assert.NoError(t, err)
+			servicePortNames := []string{}
+			for _, port := range svc.Spec.Ports {
+				servicePortNames = append(servicePortNames, port.Name)
+			}
+			expectedPortNames := []string{"mongot-grpc", "metrics", "healthcheck"}
+			if tc.withWireproto {
+				expectedPortNames = append(expectedPortNames, "mongot-wireproto")
+			}
+			assert.ElementsMatch(t, expectedPortNames, servicePortNames)
+
+			cm := &corev1.ConfigMap{}
+			err = c.Get(ctx, search.MongotConfigConfigMapNamespacedName(), cm)
+			assert.NoError(t, err)
+			expectedConfig := buildExpectedMongotConfig(search, mdbc)
+			configYaml, err := yaml.Marshal(expectedConfig)
+			assert.NoError(t, err)
+			assert.Equal(t, string(configYaml), cm.Data[searchcontroller.MongotConfigFilename])
+
+			sts := &appsv1.StatefulSet{}
+			err = c.Get(ctx, search.StatefulSetNamespacedName(), sts)
+			assert.NoError(t, err)
+		})
+	}
 }
 
 func checkSearchReconcileFailed(

controllers/searchcontroller/mongodbsearch_reconcile_helper.go

@@ -317,12 +317,14 @@ func buildSearchHeadlessService(search *searchv1.MongoDBSearch) corev1.Service {
 		SetPublishNotReadyAddresses(true).
 		SetOwnerReferences(search.GetOwnerReferences())
 
-	serviceBuilder.AddPort(&corev1.ServicePort{
-		Name:       "mongot-wireproto",
-		Protocol:   corev1.ProtocolTCP,
-		Port:       search.GetMongotWireprotoPort(),
-		TargetPort: intstr.FromInt32(search.GetMongotWireprotoPort()),
-	})
+	if search.IsWireprotoForced() {
+		serviceBuilder.AddPort(&corev1.ServicePort{
+			Name:       "mongot-wireproto",
+			Protocol:   corev1.ProtocolTCP,
+			Port:       search.GetMongotWireprotoPort(),
+			TargetPort: intstr.FromInt32(search.GetMongotWireprotoPort()),
+		})
+	}
 
 	serviceBuilder.AddPort(&corev1.ServicePort{
 		Name:       "mongot-grpc",

controllers/searchcontroller/mongodbsearch_reconcile_helper_test.go

@@ -1,6 +1,8 @@
 package searchcontroller
 
 import (
+	"fmt"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -146,3 +148,58 @@ func TestNeedsSearchCoordinatorRolePolyfill(t *testing.T) {
 		})
 	}
 }
+
+func TestGetMongodConfigParameters_TransportAndPorts(t *testing.T) {
+	cases := []struct {
+		name            string
+		withWireproto   bool
+		expectedUseGrpc bool
+	}{
+		{
+			name:            "grpc only (default)",
+			withWireproto:   false,
+			expectedUseGrpc: true,
+		},
+		{
+			name:            "grpc + wireproto via annotation",
+			withWireproto:   true,
+			expectedUseGrpc: false,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			search := &searchv1.MongoDBSearch{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-mongodb-search",
+					Namespace: "test",
+				},
+			}
+			if tc.withWireproto {
+				search.Annotations = map[string]string{searchv1.ForceWireprotoTransportAnnotation: "true"}
+			}
+
+			clusterDomain := "cluster.local"
+			params := GetMongodConfigParameters(search, clusterDomain)
+
+			setParams := params["setParameter"].(map[string]any)
+
+			useGrpc := setParams["useGrpcForSearch"].(bool)
+			assert.Equal(t, tc.expectedUseGrpc, useGrpc)
+
+			expectedPort := search.GetMongotGrpcPort()
+			if tc.withWireproto {
+				expectedPort = search.GetMongotWireprotoPort()
+			}
+			expectedPrefix := fmt.Sprintf("%s.%s.svc.%s", search.Name+"-search-svc", search.Namespace, clusterDomain)
+			expectedSuffix := fmt.Sprintf(":%d", expectedPort)
+
+			for _, key := range []string{"mongotHost", "searchIndexManagementHostAndPort"} {
+				value := setParams[key].(string)
+				if !strings.HasPrefix(value, expectedPrefix) || !strings.HasSuffix(value, expectedSuffix) {
+					t.Fatalf("%s mismatch: expected prefix %q and suffix %q, got %q", key, expectedPrefix, expectedSuffix, value)
+				}
+			}
+		})
+	}
+}