From f633696634f9292769511083a65f4becea9c52c1 Mon Sep 17 00:00:00 2001
From: Vincent Bernat <vincent@bernat.ch>
Date: Tue, 2 Dec 2025 10:11:43 +0100
Subject: [PATCH] inlet/kafka: also fix test for random load balancing

We expect a few messages in each partition.
---
 inlet/kafka/root_test.go | 11 ++++++-----
 inlet/kafka/tests.go     |  4 ++++
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/inlet/kafka/root_test.go b/inlet/kafka/root_test.go
index 94daf28c..9b0f32aa 100644
--- a/inlet/kafka/root_test.go
+++ b/inlet/kafka/root_test.go
@@ -159,14 +159,15 @@ func TestLoadBalancingAlgorithm(t *testing.T) {
 			}
 			wg.Wait()
 
-			expected := make(map[int32]int, len(messages))
+			expected := make(map[int32]int, DefaultMockKafkaNumPartitions)
 			switch algo {
 			case LoadBalanceRandom:
-				for p, count := range messages {
-					if count > total/len(messages)*2/10 && count < total/len(messages)*18/10 {
-						expected[p] = count
+				for p := range DefaultMockKafkaNumPartitions {
+					p := int32(p)
+					if messages[p] > total/DefaultMockKafkaNumPartitions*2/10 {
+						expected[p] = messages[p]
 					} else {
-						expected[p] = total / len(messages)
+						expected[p] = total / DefaultMockKafkaNumPartitions
 					}
 				}
 			case LoadBalanceByExporter:
diff --git a/inlet/kafka/tests.go b/inlet/kafka/tests.go
index d0119076..d6561d22 100644
--- a/inlet/kafka/tests.go
+++ b/inlet/kafka/tests.go
@@ -16,6 +16,9 @@ import (
 	"akvorado/common/reporter"
 )
 
+// DefaultMockKafkaNumPartitions is the default number of partitions for the mock Kafka.
+const DefaultMockKafkaNumPartitions = 8
+
 // NewMock creates a new Kafka component with a mocked Kafka. It will
 // panic if it cannot be started.
 func NewMock(t *testing.T, r *reporter.Reporter, configuration Configuration) (*Component, *kfake.Cluster) {
@@ -24,6 +27,7 @@ func NewMock(t *testing.T, r *reporter.Reporter, configuration Configuration) (*
 	cluster, err := kfake.NewCluster(
 		kfake.NumBrokers(1),
 		kfake.AllowAutoTopicCreation(),
+		kfake.DefaultNumPartitions(DefaultMockKafkaNumPartitions),
 		kfake.WithLogger(kafka.NewLogger(r)),
 	)
 	if err != nil {