init ingestion

ingestion/bulker/internal/batcher/batcher_test.go

@@ -0,0 +1,56 @@
+package batcher
+
+import (
+	"context"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"go.uber.org/zap"
+
+	"github.com/dbiz/cdp/ingestion/bulker/internal/model"
+)
+
+func TestBatcher_FlushesOnSizeCap(t *testing.T) {
+	var flushed int32
+	flush := func(_ context.Context, evs []*model.IngestedEvent) error {
+		atomic.AddInt32(&flushed, int32(len(evs)))
+		return nil
+	}
+	b := New(3, time.Hour, flush, zap.NewNop())
+	for i := 0; i < 3; i++ {
+		_ = b.Add(context.Background(), &model.IngestedEvent{MessageID: "x"})
+	}
+	assert.Equal(t, int32(3), atomic.LoadInt32(&flushed))
+}
+
+func TestBatcher_FlushNow_NoOpOnEmpty(t *testing.T) {
+	var called int32
+	flush := func(_ context.Context, _ []*model.IngestedEvent) error {
+		atomic.AddInt32(&called, 1)
+		return nil
+	}
+	b := New(10, time.Hour, flush, zap.NewNop())
+	_ = b.FlushNow(context.Background())
+	assert.Equal(t, int32(0), atomic.LoadInt32(&called))
+}
+
+func TestBatcher_FlushesOnTimer(t *testing.T) {
+	var flushed int32
+	flush := func(_ context.Context, evs []*model.IngestedEvent) error {
+		atomic.AddInt32(&flushed, int32(len(evs)))
+		return nil
+	}
+	b := New(1000, 50*time.Millisecond, flush, zap.NewNop())
+	ctx, cancel := context.WithCancel(context.Background())
+	go b.Run(ctx)
+
+	_ = b.Add(context.Background(), &model.IngestedEvent{MessageID: "a"})
+	_ = b.Add(context.Background(), &model.IngestedEvent{MessageID: "b"})
+
+	time.Sleep(120 * time.Millisecond)
+	cancel()
+	time.Sleep(10 * time.Millisecond)
+	assert.Equal(t, int32(2), atomic.LoadInt32(&flushed))
+}