57 lines
1.5 KiB
Go
57 lines
1.5 KiB
Go
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))
|
|
}
|