adding test with interval

projectdiscovery · Jun 22, 2023 · 6c3f22b · 6c3f22b
1 parent 60f54e5
commit 6c3f22b
Showing 1 changed file with 37 additions and 1 deletion.
diff --git a/batcher/batcher_test.go b/batcher/batcher_test.go
@@ -2,11 +2,12 @@ package batcher
 
 import (
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 )
 
-func TestBatcher(t *testing.T) {
+func TestBatcherStandard(t *testing.T) {
 	var (
 		batchSize        = 100
 		wanted           = 100000
@@ -38,3 +39,38 @@ func TestBatcher(t *testing.T) {
 	require.Equal(t, wanted, got)
 	require.True(t, minWantedBatches <= gotBatches)
 }
+
+func TestBatcherWithInterval(t *testing.T) {
+	var (
+		batchSize        = 200
+		wanted           = 1000
+		minWantedBatches = 10
+		got              int
+		gotBatches       int
+	)
+	callback := func(t []int) {
+		gotBatches++
+		for range t {
+			got++
+		}
+	}
+	bat := New[int](
+		WithMaxCapacity[int](batchSize),
+		WithFlushCallback[int](callback),
+		WithFlushInterval[int](10*time.Millisecond),
+	)
+
+	bat.Run()
+
+	for i := 0; i < wanted; i++ {
+		time.Sleep(2 * time.Millisecond)
+		bat.Append(i)
+	}
+
+	bat.Stop()
+
+	bat.WaitDone()
+
+	require.Equal(t, wanted, got)
+	require.True(t, minWantedBatches <= gotBatches)
+}