fix(test): guard counter in parallel

Signed-off-by: Nguyen Khac Thanh <[email protected]>

notify/kafka/kafka.go

@@ -34,7 +34,7 @@ import (
 type Notifier struct {
 	conf           *config.KafkaConfig
 	logger         *slog.Logger
-	partition      int
+	Partition      int
 	partitionMutex sync.Mutex
 	sendFunc       func(ctx context.Context, msgs ...ckafka.Message) error
 }
@@ -80,9 +80,9 @@ func New(c *config.KafkaConfig, l *slog.Logger, sendFunc *func(ctx context.Conte
 	}
 
 	if c.NumberOfPartition > 0 {
-		n.partition = c.NumberOfPartition
+		n.Partition = c.NumberOfPartition
 	} else {
-		n.partition = 1
+		n.Partition = 1
 	}
 
 	n.sendFunc = func(ctx context.Context, msgs ...ckafka.Message) error {
@@ -94,14 +94,16 @@ func New(c *config.KafkaConfig, l *slog.Logger, sendFunc *func(ctx context.Conte
 
 // GetPartitionIndex returns the current partition index.
 func (n *Notifier) GetPartitionIndex() int {
-	return n.partition
+	n.partitionMutex.Lock()
+	defer n.partitionMutex.Unlock()
+	return n.Partition
 }
 
 // NextPartition returns the next partition index.
 func (n *Notifier) NextPartition() {
 	n.partitionMutex.Lock()
-	n.partition = (n.partition + 1) % n.conf.NumberOfPartition
-	n.partitionMutex.Unlock()
+	defer n.partitionMutex.Unlock()
+	n.Partition = (n.Partition + 1) % n.conf.NumberOfPartition
 }
 
 // Notify implements the Notifier interface.

notify/kafka/kafka_test.go

@@ -28,9 +28,7 @@ import (
 )
 
 func TestKafkaNotify(t *testing.T) {
-	var counter int
 	sendFunc := func(ctx context.Context, msgs ...ckafka.Message) error {
-		counter++
 		return nil
 	}
 
@@ -47,13 +45,18 @@ func TestKafkaNotify(t *testing.T) {
 
 	notifier.Notify(context.Background(), &types.Alert{})
 
-	require.Equal(t, 1, counter)
+	require.Equal(t, 0, notifier.Partition)
 }
 
 func TestKafkaNotifyRoundRobin(t *testing.T) {
-	var counter int
+	var (
+		counter      int
+		counterMutex sync.Mutex
+	)
 	partitions := 2
 	sendFunc := func(ctx context.Context, msgs ...ckafka.Message) error {
+		counterMutex.Lock()
+		defer counterMutex.Unlock()
 		require.Equal(t, counter%partitions, msgs[0].Partition)
 		counter++
 		return nil
@@ -73,10 +76,6 @@ func TestKafkaNotifyRoundRobin(t *testing.T) {
 
 	var wg sync.WaitGroup
 
-	notifier.Notify(context.Background(), &types.Alert{})
-	notifier.Notify(context.Background(), &types.Alert{})
-	notifier.Notify(context.Background(), &types.Alert{})
-
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
@@ -89,8 +88,7 @@ func TestKafkaNotifyRoundRobin(t *testing.T) {
 		notifier.Notify(context.Background(), &types.Alert{})
 	}()
 
-	notifier.Notify(context.Background(), &types.Alert{})
-	notifier.Notify(context.Background(), &types.Alert{})
-
 	wg.Wait()
+
+	require.Equal(t, partitions, counter)
 }