fix: unit test failure

Author: siddhanta-rath

.env.sample

@@ -27,7 +27,9 @@ PUBLISHER_KAFKA_CLIENT_STATISTICS_INTERVAL_MS=5000
 PUBLISHER_KAFKA_CLIENT_QUEUE_BUFFERING_MAX_MESSAGES=100000
 PUBLISHER_KAFKA_FLUSH_INTERVAL_MS=1000
 
-c
+STATS_TOPIC_NAME="clickstream-total-event-log"
+STATS_FLUSH_INTERVAL_IN_SEC=10
+STATS_CHANNEL_SIZE=1000
 
 METRIC_STATSD_ADDRESS=":8125"
 METRIC_STATSD_FLUSH_PERIOD_MS=100

stats/totalevent_test.go

@@ -12,7 +12,8 @@ import (
 func TestFlushTotalEventStat(t *testing.T) {
 	// Arrange
 	eventCh := make(chan int32, 10)
-	mockKP := &mockKafkaProducer{}
+	producedCh := make(chan *pb.TotalEventCountMessage, 10) // <- initialize channel
+	mockKP := &mockKafkaProducer{producedCh: producedCh}
 	flushInterval := 1 * time.Millisecond
 	topicName := "test-topic"
 
@@ -27,26 +28,24 @@ func TestFlushTotalEventStat(t *testing.T) {
 	time.Sleep(2 * flushInterval)
 
 	// Assert
-	if len(mockKP.produced) == 0 {
+	select {
+	case msg := <-producedCh:
+		if msg.EventCount != 5 {
+			t.Errorf("expected EventCount=5, got %d", msg.EventCount)
+		}
+		if msg.EventTimestamp.AsTime().IsZero() {
+			t.Error("expected non-zero EventTimestamp")
+		}
+	default:
 		t.Error("expected ProduceEventStat to be called, got 0")
 	}
-
-	result := mockKP.produced[0].EventCount
-	expected := int32(5)
-	if result != expected {
-		t.Errorf("expected EventCount=%d, got %d", expected, result)
-	}
-
-	// Also check timestamp is set
-	if mockKP.produced[0].EventTimestamp.AsTime().IsZero() {
-		t.Error("expected non-zero EventTimestamp")
-	}
 }
 
 // Optional: add a timeout test to ensure aggregation resets
 func TestFlushResetsAfterPublish(t *testing.T) {
 	eventCh := make(chan int32, 10)
-	mockKP := &mockKafkaProducer{}
+	producedCh := make(chan *pb.TotalEventCountMessage, 10)
+	mockKP := &mockKafkaProducer{producedCh: producedCh}
 	flushInterval := 1 * time.Millisecond
 	ts := stats.CreateTotalEventStat(mockKP, flushInterval, "topic", eventCh)
 
@@ -60,28 +59,37 @@ func TestFlushResetsAfterPublish(t *testing.T) {
 	eventCh <- 4
 	time.Sleep(2 * flushInterval)
 
-	if len(mockKP.produced) < 2 {
-		t.Fatalf("expected at least 2 ProduceEventStat calls, got %d", len(mockKP.produced))
+	// Assert first batch
+	select {
+	case msg := <-producedCh:
+		if msg.EventCount != 1 {
+			t.Errorf("expected first flush=1, got %d", msg.EventCount)
+		}
+	default:
+		t.Fatal("expected first ProduceEventStat call")
 	}
 
-	if mockKP.produced[0].EventCount != 1 {
-		t.Errorf("expected first flush=1, got %d", mockKP.produced[0].EventCount)
-	}
-	if mockKP.produced[1].EventCount != 4 {
-		t.Errorf("expected second flush=4, got %d", mockKP.produced[1].EventCount)
+	// Assert second batch
+	select {
+	case msg := <-producedCh:
+		if msg.EventCount != 4 {
+			t.Errorf("expected second flush=4, got %d", msg.EventCount)
+		}
+	default:
+		t.Fatal("expected second ProduceEventStat call")
 	}
 }
 
 // mockKafkaProducer implements publisher.KafkaProducer
 type mockKafkaProducer struct {
-	produced []*pb.TotalEventCountMessage
+	producedCh chan *pb.TotalEventCountMessage
 }
 
 func (m *mockKafkaProducer) ProduceBulk(_ []*pb.Event, _ string, _ chan kafka.Event) error {
 	return nil
 }
 
 func (m *mockKafkaProducer) ProduceEventStat(_ string, event *pb.TotalEventCountMessage) error {
-	m.produced = append(m.produced, event)
+	m.producedCh <- event
 	return nil
 }