From 013c5722b6a0d2b30dd49714120bab652f3536f1 Mon Sep 17 00:00:00 2001
From: Amit Slavin <108348428+amitslavin@users.noreply.github.com>
Date: Thu, 5 Dec 2024 17:00:27 +0200
Subject: [PATCH] [USM] Fix TestInvalidBatchCountMetric flakiness (#31772)

---
 pkg/network/protocols/events/consumer_test.go | 71 ++++++++++++-------
 1 file changed, 45 insertions(+), 26 deletions(-)

diff --git a/pkg/network/protocols/events/consumer_test.go b/pkg/network/protocols/events/consumer_test.go
index 2b1bba26b97ad..5e9864b5a342a 100644
--- a/pkg/network/protocols/events/consumer_test.go
+++ b/pkg/network/protocols/events/consumer_test.go
@@ -18,6 +18,8 @@ import (
 
 	manager "github.com/DataDog/ebpf-manager"
 	"github.com/cilium/ebpf"
+	"github.com/cilium/ebpf/features"
+	"github.com/cilium/ebpf/perf"
 	"github.com/cilium/ebpf/ringbuf"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -78,6 +80,33 @@ func TestConsumer(t *testing.T) {
 	}
 }
 
+func TestInvalidBatchCountMetric(t *testing.T) {
+	kversion, err := kernel.HostVersion()
+	require.NoError(t, err)
+	if minVersion := kernel.VersionCode(4, 14, 0); kversion < minVersion {
+		t.Skipf("package not supported by kernels < %s", minVersion)
+	}
+
+	c := config.New()
+	program, err := newEBPFProgram(c)
+	require.NoError(t, err)
+	t.Cleanup(func() { program.Stop(manager.CleanAll) })
+
+	consumer, err := NewConsumer("test", program, func([]uint64) {})
+	require.NoError(t, err)
+
+	// We are creating a raw sample with a data length of 4, which is smaller than sizeOfBatch
+	// and would be considered an invalid batch.
+	recordSample(c, consumer, []byte("test"))
+
+	consumer.Start()
+	t.Cleanup(func() { consumer.Stop() })
+	require.Eventually(t, func() bool {
+		// Wait for the consumer to process the invalid batch.
+		return consumer.invalidBatchCount.Get() == 1
+	}, 5*time.Second, 100*time.Millisecond)
+}
+
 type eventGenerator struct {
 	// map used for coordinating test with eBPF program space
 	testMap *ebpf.Map
@@ -86,6 +115,22 @@ type eventGenerator struct {
 	testFile *os.File
 }
 
+// recordSample records a sample using the consumer handler.
+func recordSample(c *config.Config, consumer *Consumer[uint64], sampleData []byte) {
+	// Ring buffers require kernel version 5.8.0 or higher, therefore, the handler is chosen based on the kernel version.
+	if c.EnableUSMRingBuffers && features.HaveMapType(ebpf.RingBuf) == nil {
+		handler := consumer.handler.(*ddebpf.RingBufferHandler)
+		handler.RecordHandler(&ringbuf.Record{
+			RawSample: sampleData,
+		}, nil, nil)
+	} else {
+		handler := consumer.handler.(*ddebpf.PerfHandler)
+		handler.RecordHandler(&perf.Record{
+			RawSample: sampleData,
+		}, nil, nil)
+	}
+}
+
 func newEventGenerator(program *manager.Manager, t *testing.T) *eventGenerator {
 	m, _, _ := program.GetMap("test")
 	require.NotNilf(t, m, "couldn't find test map")
@@ -171,29 +216,3 @@ func newEBPFProgram(c *config.Config) (*manager.Manager, error) {
 
 	return m, nil
 }
-
-func TestInvalidBatchCountMetric(t *testing.T) {
-	kversion, err := kernel.HostVersion()
-	require.NoError(t, err)
-	if minVersion := kernel.VersionCode(4, 14, 0); kversion < minVersion {
-		t.Skipf("package not supported by kernels < %s", minVersion)
-	}
-
-	program, err := newEBPFProgram(config.New())
-	require.NoError(t, err)
-
-	ringBufferHandler := ddebpf.NewRingBufferHandler(1)
-	ringBufferHandler.RecordHandler(&ringbuf.Record{
-		RawSample: []byte("test"),
-	}, nil, nil)
-
-	consumer, err := NewConsumer("test", program, func(_ []uint64) {})
-	require.NoError(t, err)
-	consumer.handler = ringBufferHandler
-
-	consumer.Start()
-	program.Stop(manager.CleanAll)
-	consumer.Stop()
-
-	require.Equalf(t, int(consumer.invalidBatchCount.Get()), 1, "invalidBatchCount should be greater than 0")
-}