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,142 @@
+// 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.Error(t, err)
+
+	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)
+	authConfig := &AuthenticationConfig{}
+
+	// Helper function to parse flags and initialize authConfig
+	parseFlagsAndInit := func(authType string) {
+		command.ParseFlags([]string{
+			"--kafka.auth.authentication=" + authType,
+			"--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.kerberos.use-keytab=false",
+		})
+		authConfig = &AuthenticationConfig{}
+		err := authConfig.InitFromViper(configPrefix, v)
+		require.NoError(t, err)
+	}
+
+	// Test with invalid authentication method
+	parseFlagsAndInit("fail")
+	require.Error(t, authConfig.SetConfiguration(saramaConfig, logger), "Unknown/Unsupported authentication method fail to kafka cluster")
+
+	// Test with kerberos
+	parseFlagsAndInit("kerberos")
+	require.NoError(t, authConfig.SetConfiguration(saramaConfig, logger))
+
+	// Test all plaintext options
+	parseFlagsAndInit("plaintext")
+	testPlaintext(v, t, configPrefix, logger, "SCRAM-SHA-256", saramaConfig)
+	testPlaintext(v, t, configPrefix, logger, "SCRAM-SHA-512", saramaConfig)
+	testPlaintext(v, t, configPrefix, logger, "PLAIN", saramaConfig)
+
+	// Test with no authentication
+	parseFlagsAndInit(" ")
+	require.NoError(t, authConfig.SetConfiguration(saramaConfig, logger))
+
+	// Test with tls
+	parseFlagsAndInit("tls")
+	require.NoError(t, authConfig.SetConfiguration(saramaConfig, logger))
+	defer authConfig.TLS.Close()
+	// test tls_fail
+	authConfig.TLS.CipherSuites = []string{"fail"}
+	require.Error(t, authConfig.SetConfiguration(saramaConfig, logger))
+}

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)
+			}
+		})
+	}
+}

pkg/kafka/auth/plaintext_test.go

@@ -0,0 +1,105 @@
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package auth
+
+import (
+	"fmt"
+	"strings"
+	"testing"
+
+	"github.com/Shopify/sarama"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/xdg-go/scram"
+)
+
+func TestScramClient(t *testing.T) {
+	scramClientFunc := clientGenFunc(scram.SHA256)
+	client := scramClientFunc().(*scramClient)
+
+	err := client.Begin("testUser", "testPassword", "testAuthzID")
+	require.NoError(t, err, "Begin should not return an error")
+	assert.NotNil(t, client.Client, "Client should be initialized")
+	assert.NotNil(t, client.ClientConversation, "ClientConversation should be initialized")
+
+	step, err := client.Step("testChallenge")
+	require.NoError(t, err, "Step should not return an error")
+	require.NotEmpty(t, step, "Step should return a non-empty response")
+
+	done := client.Done()
+	assert.False(t, done, "Done should return false initially")
+}
+
+func TestSetPlainTextConfiguration(t *testing.T) {
+	tests := []struct {
+		name              string
+		config            PlainTextConfig
+		expectedError     error
+		expectedMechanism sarama.SASLMechanism
+	}{
+		{
+			name: "SCRAM-SHA-256",
+			config: PlainTextConfig{
+				Username:  "username",
+				Password:  "password",
+				Mechanism: "SCRAM-SHA-256",
+			},
+			expectedError:     nil,
+			expectedMechanism: sarama.SASLTypeSCRAMSHA256,
+		},
+		{
+			name: "SCRAM-SHA-512",
+			config: PlainTextConfig{
+				Username:  "username",
+				Password:  "password",
+				Mechanism: "SCRAM-SHA-512",
+			},
+			expectedError:     nil,
+			expectedMechanism: sarama.SASLTypeSCRAMSHA512,
+		},
+		{
+			name: "PLAIN",
+			config: PlainTextConfig{
+				Username:  "username",
+				Password:  "password",
+				Mechanism: "PLAIN",
+			},
+			expectedError:     nil,
+			expectedMechanism: sarama.SASLTypePlaintext,
+		},
+		{
+			name: "Invalid Mechanism",
+			config: PlainTextConfig{
+				Username:  "username",
+				Password:  "password",
+				Mechanism: "INVALID",
+			},
+			expectedError: fmt.Errorf("config plaintext.mechanism error: %s, only support 'SCRAM-SHA-256' or 'SCRAM-SHA-512' or 'PLAIN'", "INVALID"),
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			saramaConfig := sarama.NewConfig()
+
+			err := setPlainTextConfiguration(&tt.config, saramaConfig)
+
+			if tt.expectedError != nil {
+				assert.EqualError(t, err, tt.expectedError.Error())
+			} else {
+				require.NoError(t, err)
+				assert.True(t, saramaConfig.Net.SASL.Enable)
+				assert.Equal(t, tt.config.Username, saramaConfig.Net.SASL.User)
+				assert.Equal(t, tt.config.Password, saramaConfig.Net.SASL.Password)
+				assert.Equal(t, tt.expectedMechanism, saramaConfig.Net.SASL.Mechanism)
+
+				if strings.HasPrefix(tt.config.Mechanism, "SCRAM-SHA-") {
+					assert.NotNil(t, saramaConfig.Net.SASL.SCRAMClientGeneratorFunc)
+				} else {
+					assert.Nil(t, saramaConfig.Net.SASL.SCRAMClientGeneratorFunc)
+				}
+			}
+		})
+	}
+}