复制项目

pkg/tools/batcher/batcher_test.go

@@ -0,0 +1,66 @@
+package batcher
+
+import (
+	"context"
+	"fmt"
+	"github.com/openimsdk/tools/utils/stringutil"
+	"testing"
+	"time"
+)
+
+func TestBatcher(t *testing.T) {
+	config := Config{
+		size:     1000,
+		buffer:   10,
+		worker:   10,
+		interval: 5 * time.Millisecond,
+	}
+
+	b := New[string](
+		WithSize(config.size),
+		WithBuffer(config.buffer),
+		WithWorker(config.worker),
+		WithInterval(config.interval),
+		WithSyncWait(true),
+	)
+
+	// Mock Do function to simply print values for demonstration
+	b.Do = func(ctx context.Context, channelID int, vals *Msg[string]) {
+		t.Logf("Channel %d Processed batch: %v", channelID, vals)
+	}
+	b.OnComplete = func(lastMessage *string, totalCount int) {
+		t.Logf("Completed processing with last message: %v, total count: %d", *lastMessage, totalCount)
+	}
+	b.Sharding = func(key string) int {
+		hashCode := stringutil.GetHashCode(key)
+		return int(hashCode) % config.worker
+	}
+	b.Key = func(data *string) string {
+		return *data
+	}
+
+	err := b.Start()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Test normal data processing
+	for i := 0; i < 10000; i++ {
+		data := "data" + fmt.Sprintf("%d", i)
+		if err := b.Put(context.Background(), &data); err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	time.Sleep(time.Duration(1) * time.Second)
+	start := time.Now()
+	// Wait for all processing to finish
+	b.Close()
+
+	elapsed := time.Since(start)
+	t.Logf("Close took %s", elapsed)
+
+	if len(b.data) != 0 {
+		t.Error("Data channel should be empty after closing")
+	}
+}