feat: support push with key

Author: kevwan

go.mod

@@ -42,6 +42,7 @@ require (
 	github.com/prometheus/procfs v0.12.0 // indirect
 	github.com/redis/go-redis/v9 v9.5.3 // indirect
 	github.com/spaolacci/murmur3 v1.1.0 // indirect
+	github.com/stretchr/objx v0.5.2 // indirect
 	go.opentelemetry.io/otel/exporters/jaeger v1.17.0 // indirect
 	go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.19.0 // indirect
 	go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.19.0 // indirect

kq/pusher.go

@@ -16,11 +16,16 @@ type (
 	PushOption func(options *pushOptions)
 
 	Pusher struct {
-		producer *kafka.Writer
 		topic    string
+		producer kafkaWriter
 		executor *executors.ChunkExecutor
 	}
 
+	kafkaWriter interface {
+		Close() error
+		WriteMessages(ctx context.Context, msgs ...kafka.Message) error
+	}
+
 	pushOptions struct {
 		// kafka.Writer options
 		allowAutoTopicCreation bool

kq/pusher_test.go

@@ -0,0 +1,141 @@
+package kq
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/segmentio/kafka-go"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+)
+
+// MockKafkaWriter is a mock for kafka.Writer
+type MockKafkaWriter struct {
+	mock.Mock
+}
+
+func (m *MockKafkaWriter) WriteMessages(ctx context.Context, msgs ...kafka.Message) error {
+	args := m.Called(ctx, msgs)
+	return args.Error(0)
+}
+
+func (m *MockKafkaWriter) Close() error {
+	args := m.Called()
+	return args.Error(0)
+}
+
+func TestNewPusher(t *testing.T) {
+	addrs := []string{"localhost:9092"}
+	topic := "test-topic"
+
+	t.Run("DefaultOptions", func(t *testing.T) {
+		pusher := NewPusher(addrs, topic)
+		assert.NotNil(t, pusher)
+		assert.NotNil(t, pusher.producer)
+		assert.Equal(t, topic, pusher.topic)
+		assert.NotNil(t, pusher.executor)
+	})
+
+	t.Run("WithSyncPush", func(t *testing.T) {
+		pusher := NewPusher(addrs, topic, WithSyncPush())
+		assert.NotNil(t, pusher)
+		assert.NotNil(t, pusher.producer)
+		assert.Equal(t, topic, pusher.topic)
+		assert.Nil(t, pusher.executor)
+	})
+
+	t.Run("WithChunkSize", func(t *testing.T) {
+		pusher := NewPusher(addrs, topic, WithChunkSize(100))
+		assert.NotNil(t, pusher)
+		assert.NotNil(t, pusher.executor)
+	})
+
+	t.Run("WithFlushInterval", func(t *testing.T) {
+		pusher := NewPusher(addrs, topic, WithFlushInterval(time.Second))
+		assert.NotNil(t, pusher)
+		assert.NotNil(t, pusher.executor)
+	})
+
+	t.Run("WithAllowAutoTopicCreation", func(t *testing.T) {
+		pusher := NewPusher(addrs, topic, WithAllowAutoTopicCreation())
+		assert.NotNil(t, pusher)
+		assert.True(t, pusher.producer.(*kafka.Writer).AllowAutoTopicCreation)
+	})
+}
+
+func TestPusher_Close(t *testing.T) {
+	mockWriter := new(MockKafkaWriter)
+	pusher := &Pusher{
+		producer: mockWriter,
+	}
+
+	mockWriter.On("Close").Return(nil)
+
+	err := pusher.Close()
+	assert.NoError(t, err)
+	mockWriter.AssertExpectations(t)
+}
+
+func TestPusher_Name(t *testing.T) {
+	topic := "test-topic"
+	pusher := &Pusher{topic: topic}
+
+	assert.Equal(t, topic, pusher.Name())
+}
+
+func TestPusher_Push(t *testing.T) {
+	mockWriter := new(MockKafkaWriter)
+	pusher := &Pusher{
+		producer: mockWriter,
+		topic:    "test-topic",
+	}
+
+	ctx := context.Background()
+	value := "test-value"
+
+	mockWriter.On("WriteMessages", mock.Anything, mock.AnythingOfType("[]kafka.Message")).Return(nil)
+
+	err := pusher.Push(ctx, value)
+	assert.NoError(t, err)
+	mockWriter.AssertExpectations(t)
+}
+
+func TestPusher_PushWithKey(t *testing.T) {
+	mockWriter := new(MockKafkaWriter)
+	pusher := &Pusher{
+		producer: mockWriter,
+		topic:    "test-topic",
+	}
+
+	ctx := context.Background()
+	key := "test-key"
+	value := "test-value"
+
+	mockWriter.On("WriteMessages", mock.Anything, mock.AnythingOfType("[]kafka.Message")).Return(nil)
+
+	err := pusher.PushWithKey(ctx, key, value)
+	assert.NoError(t, err)
+	mockWriter.AssertExpectations(t)
+}
+
+func TestPusher_PushWithKey_Error(t *testing.T) {
+	mockWriter := new(MockKafkaWriter)
+	pusher := &Pusher{
+		producer: mockWriter,
+		topic:    "test-topic",
+	}
+
+	ctx := context.Background()
+	key := "test-key"
+	value := "test-value"
+
+	expectedError := errors.New("write error")
+	mockWriter.On("WriteMessages", mock.Anything, mock.AnythingOfType("[]kafka.Message")).Return(expectedError)
+
+	err := pusher.PushWithKey(ctx, key, value)
+	assert.Error(t, err)
+	assert.Equal(t, expectedError, err)
+	mockWriter.AssertExpectations(t)
+}