Increase codecov of pkg/kafka

internal/tracegen/worker_test.go

@@ -16,30 +16,47 @@ import (
 )
 
 func Test_SimulateTraces(t *testing.T) {
-	logger, buf := testutils.NewLogger()
-	tp := sdktrace.NewTracerProvider()
-	tracers := []trace.Tracer{tp.Tracer("stdout")}
-	wg := sync.WaitGroup{}
-	wg.Add(1)
-	var running uint32 = 1
-
-	worker := &worker{
-		logger:  logger,
-		tracers: tracers,
-		wg:      &wg,
-		id:      7,
-		running: &running,
-		Config: Config{
-			Traces:   7,
-			Duration: time.Second,
-			Pause:    time.Second,
-			Service:  "stdout",
-			Debug:    true,
-			Firehose: true,
+	tests := []struct {
+		name  string
+		pause time.Duration
+	}{
+		{
+			name:  "no pause",
+			pause: 0,
+		},
+		{
+			name:  "with pause",
+			pause: time.Second,
 		},
 	}
-	expectedOutput := `{"level":"info","msg":"Worker 7 generated 7 traces"}` + "\n"
 
-	worker.simulateTraces()
-	assert.Equal(t, expectedOutput, buf.String())
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			logger, buf := testutils.NewLogger()
+			tp := sdktrace.NewTracerProvider()
+			tracers := []trace.Tracer{tp.Tracer("stdout")}
+			wg := sync.WaitGroup{}
+			wg.Add(1)
+			var running uint32 = 1
+			worker := &worker{
+				logger:  logger,
+				tracers: tracers,
+				wg:      &wg,
+				id:      7,
+				running: &running,
+				Config: Config{
+					Traces:     7,
+					Duration:   time.Second,
+					Pause:      tt.pause,
+					Service:    "stdout",
+					Debug:      true,
+					Firehose:   true,
+					ChildSpans: 1,
+				},
+			}
+			expectedOutput := `{"level":"info","msg":"Worker 7 generated 7 traces"}` + "\n"
+			worker.simulateTraces()
+			assert.Equal(t, expectedOutput, buf.String())
+		})
+	}
 }

pkg/kafka/auth/config_test.go

@@ -0,0 +1,173 @@
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package auth
+
+import (
+	"flag"
+	"testing"
+
+	"github.com/Shopify/sarama"
+	"github.com/spf13/viper"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/zap"
+
+	"github.com/jaegertracing/jaeger/pkg/config"
+	"github.com/jaegertracing/jaeger/pkg/config/tlscfg"
+)
+
+func addFlags(flags *flag.FlagSet) {
+	configPrefix := "kafka.auth"
+	AddFlags(configPrefix, flags)
+}
+
+func Test_InitFromViper(t *testing.T) {
+	configPrefix := "kafka.auth"
+	v, command := config.Viperize(addFlags)
+	command.ParseFlags([]string{
+		"--kafka.auth.authentication=tls",
+		"--kafka.auth.kerberos.service-name=kafka",
+		"--kafka.auth.kerberos.realm=EXAMPLE.COM",
+		"--kafka.auth.kerberos.use-keytab=true",
+		"--kafka.auth.kerberos.username=user",
+		"--kafka.auth.kerberos.password=password",
+		"--kafka.auth.kerberos.config-file=/path/to/krb5.conf",
+		"--kafka.auth.kerberos.keytab-file=/path/to/keytab",
+		"--kafka.auth.kerberos.disable-fast-negotiation=true",
+		"--kafka.auth.tls.enabled=false",
+		"--kafka.auth.plaintext.username=user",
+		"--kafka.auth.plaintext.password=password",
+		"--kafka.auth.plaintext.mechanism=SCRAM-SHA-256",
+		"--kafka.auth.tls.ca=failing",
+	})
+
+	authConfig := &AuthenticationConfig{}
+	err := authConfig.InitFromViper(configPrefix, v)
+	require.EqualError(t, err, "failed to process Kafka TLS options: kafka.auth.tls.* options cannot be used when kafka.auth.tls.enabled is false")
+
+	command.ParseFlags([]string{"--kafka.auth.tls.ca="})
+	v.BindPFlags(command.Flags())
+	err = authConfig.InitFromViper(configPrefix, v)
+	require.NoError(t, err)
+
+	expectedConfig := &AuthenticationConfig{
+		Authentication: "tls",
+		Kerberos: KerberosConfig{
+			ServiceName:     "kafka",
+			Realm:           "EXAMPLE.COM",
+			UseKeyTab:       true,
+			Username:        "user",
+			Password:        "password",
+			ConfigPath:      "/path/to/krb5.conf",
+			KeyTabPath:      "/path/to/keytab",
+			DisablePAFXFast: true,
+		},
+		TLS: tlscfg.Options{
+			Enabled: true,
+		},
+		PlainText: PlainTextConfig{
+			Username:  "user",
+			Password:  "password",
+			Mechanism: "SCRAM-SHA-256",
+		},
+	}
+	assert.Equal(t, expectedConfig, authConfig)
+}
+
+// Test plaintext with different mechanisms
+func testPlaintext(v *viper.Viper, t *testing.T, configPrefix string, logger *zap.Logger, mechanism string, saramaConfig *sarama.Config) {
+	v.Set(configPrefix+plainTextPrefix+suffixPlainTextMechanism, mechanism)
+	authConfig := &AuthenticationConfig{}
+	err := authConfig.InitFromViper(configPrefix, v)
+	require.NoError(t, err)
+	require.NoError(t, authConfig.SetConfiguration(saramaConfig, logger))
+}
+
+func TestSetConfiguration(t *testing.T) {
+	logger, _ := zap.NewDevelopment()
+	saramaConfig := sarama.NewConfig()
+	configPrefix := "kafka.auth"
+	v, command := config.Viperize(addFlags)
+
+	// Table-driven test cases
+	tests := []struct {
+		name                string
+		authType            string
+		expectedError       string
+		plainTextMechanisms []string
+	}{
+		{
+			name:          "Invalid authentication method",
+			authType:      "fail",
+			expectedError: "Unknown/Unsupported authentication method fail to kafka cluster",
+		},
+		{
+			name:          "Kerberos authentication",
+			authType:      "kerberos",
+			expectedError: "",
+		},
+		{
+			name:                "Plaintext authentication with SCRAM-SHA-256",
+			authType:            "plaintext",
+			expectedError:       "",
+			plainTextMechanisms: []string{"SCRAM-SHA-256"},
+		},
+		{
+			name:                "Plaintext authentication with SCRAM-SHA-512",
+			authType:            "plaintext",
+			expectedError:       "",
+			plainTextMechanisms: []string{"SCRAM-SHA-512"},
+		},
+		{
+			name:                "Plaintext authentication with PLAIN",
+			authType:            "plaintext",
+			expectedError:       "",
+			plainTextMechanisms: []string{"PLAIN"},
+		},
+		{
+			name:          "No authentication",
+			authType:      " ",
+			expectedError: "",
+		},
+		{
+			name:          "TLS authentication",
+			authType:      "tls",
+			expectedError: "",
+		},
+		{
+			name:          "TLS authentication with invalid cipher suite",
+			authType:      "tls",
+			expectedError: "error loading tls config: failed to get cipher suite ids from cipher suite names: cipher suite fail not supported or doesn't exist",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			command.ParseFlags([]string{
+				"--kafka.auth.authentication=" + tt.authType,
+			})
+			authConfig := &AuthenticationConfig{}
+			defer authConfig.TLS.Close()
+			err := authConfig.InitFromViper(configPrefix, v)
+			require.NoError(t, err)
+
+			if tt.authType == "tls" && tt.expectedError != "" {
+				authConfig.TLS.CipherSuites = []string{"fail"}
+			}
+
+			if len(tt.plainTextMechanisms) > 0 {
+				for _, mechanism := range tt.plainTextMechanisms {
+					testPlaintext(v, t, configPrefix, logger, mechanism, saramaConfig)
+				}
+			} else {
+				err = authConfig.SetConfiguration(saramaConfig, logger)
+				if tt.expectedError != "" {
+					require.EqualError(t, err, tt.expectedError)
+				} else {
+					require.NoError(t, err)
+				}
+			}
+		})
+	}
+}

pkg/kafka/auth/kerberos_test.go

@@ -0,0 +1,68 @@
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package auth
+
+import (
+	"testing"
+
+	"github.com/Shopify/sarama"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestSetKerberosConfiguration(t *testing.T) {
+	tests := []struct {
+		name   string
+		config KerberosConfig
+	}{
+		{
+			name: "With KeyTab",
+			config: KerberosConfig{
+				ServiceName:     "service",
+				Realm:           "realm",
+				UseKeyTab:       true,
+				Username:        "username",
+				Password:        "password",
+				ConfigPath:      "/path/to/config",
+				KeyTabPath:      "/path/to/keytab",
+				DisablePAFXFast: true,
+			},
+		},
+		{
+			name: "Without KeyTab",
+			config: KerberosConfig{
+				ServiceName:     "service",
+				Realm:           "realm",
+				UseKeyTab:       false,
+				Username:        "username",
+				Password:        "password",
+				ConfigPath:      "/path/to/config",
+				DisablePAFXFast: false,
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			saramaConfig := sarama.NewConfig()
+
+			setKerberosConfiguration(&tt.config, saramaConfig)
+
+			assert.Equal(t, sarama.SASLMechanism("GSSAPI"), saramaConfig.Net.SASL.Mechanism)
+			assert.True(t, saramaConfig.Net.SASL.Enable)
+			assert.Equal(t, tt.config.Username, saramaConfig.Net.SASL.GSSAPI.Username)
+			assert.Equal(t, tt.config.Realm, saramaConfig.Net.SASL.GSSAPI.Realm)
+			assert.Equal(t, tt.config.ServiceName, saramaConfig.Net.SASL.GSSAPI.ServiceName)
+			assert.Equal(t, tt.config.DisablePAFXFast, saramaConfig.Net.SASL.GSSAPI.DisablePAFXFAST)
+			assert.Equal(t, tt.config.ConfigPath, saramaConfig.Net.SASL.GSSAPI.KerberosConfigPath)
+
+			if tt.config.UseKeyTab {
+				assert.Equal(t, tt.config.KeyTabPath, saramaConfig.Net.SASL.GSSAPI.KeyTabPath)
+				assert.Equal(t, sarama.KRB5_KEYTAB_AUTH, saramaConfig.Net.SASL.GSSAPI.AuthType)
+			} else {
+				assert.Equal(t, tt.config.Password, saramaConfig.Net.SASL.GSSAPI.Password)
+				assert.Equal(t, sarama.KRB5_USER_AUTH, saramaConfig.Net.SASL.GSSAPI.AuthType)
+			}
+		})
+	}
+}