im_server/server/kafka_runtime_batch_test.go

package main

import (
	"context"
	"errors"
	"strings"
	"testing"
	"time"

	"github.com/segmentio/kafka-go"
	"github.com/tinode/chat/server/store/types"
)

func TestConsumeKafkaMessageBatchRetriesHandlerAndCommit(t *testing.T) {
	handlerCalls := 0
	commitCalls := 0
	waitCalls := 0
	messages := []kafka.Message{{Offset: 1}, {Offset: 2}}

	err := consumeKafkaMessageBatch(context.Background(), messages,
		func(context.Context, []kafka.Message) error {
			handlerCalls++
			if handlerCalls == 1 {
				return errors.New("batch failed")
			}
			return nil
		},
		func(_ context.Context, committed ...kafka.Message) error {
			commitCalls++
			if len(committed) != len(messages) {
				t.Fatalf("committed %d messages, want %d", len(committed), len(messages))
			}
			if commitCalls == 1 {
				return errors.New("commit failed")
			}
			return nil
		},
		func(context.Context, time.Duration) error {
			waitCalls++
			return nil
		})

	if err != nil {
		t.Fatalf("consumeKafkaMessageBatch() error = %v", err)
	}
	if handlerCalls != 2 || commitCalls != 2 || waitCalls != 2 {
		t.Fatalf("calls handler=%d commit=%d wait=%d, want 2/2/2", handlerCalls, commitCalls, waitCalls)
	}
}

func TestConsumeKafkaMessageBatchDoesNotCommitWhileHandlerFails(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	handlerCalls := 0
	commitCalls := 0

	err := consumeKafkaMessageBatch(ctx, []kafka.Message{{Offset: 1}},
		func(context.Context, []kafka.Message) error {
			handlerCalls++
			return errors.New("database unavailable")
		},
		func(context.Context, ...kafka.Message) error {
			commitCalls++
			return nil
		},
		func(ctx context.Context, _ time.Duration) error {
			cancel()
			return ctx.Err()
		})

	if !errors.Is(err, context.Canceled) {
		t.Fatalf("consumeKafkaMessageBatch() error = %v, want context canceled", err)
	}
	if handlerCalls != 1 || commitCalls != 0 {
		t.Fatalf("calls handler=%d commit=%d, want 1/0", handlerCalls, commitCalls)
	}
}

func TestFetchKafkaMessageBatchWithRetainsPartialBatchOnFetchError(t *testing.T) {
	fetchCalls := 0
	fetchErr := errors.New("broker connection reset")
	messages, err := fetchKafkaMessageBatchWith(context.Background(), 10, time.Second,
		func(context.Context) (kafka.Message, error) {
			fetchCalls++
			if fetchCalls == 1 {
				return kafka.Message{Partition: 2, Offset: 41}, nil
			}
			return kafka.Message{}, fetchErr
		})

	if !errors.Is(err, fetchErr) {
		t.Fatalf("fetchKafkaMessageBatchWith() error = %v, want %v", err, fetchErr)
	}
	if len(messages) != 1 || messages[0].Partition != 2 || messages[0].Offset != 41 {
		t.Fatalf("fetchKafkaMessageBatchWith() messages = %#v, want retained offset 41", messages)
	}
}

func TestFetchKafkaMessageBatchWithReturnsInitialFetchErrorWithoutMessages(t *testing.T) {
	fetchErr := errors.New("broker unavailable")
	messages, err := fetchKafkaMessageBatchWith(context.Background(), 10, time.Second,
		func(context.Context) (kafka.Message, error) {
			return kafka.Message{}, fetchErr
		})

	if !errors.Is(err, fetchErr) {
		t.Fatalf("fetchKafkaMessageBatchWith() error = %v, want %v", err, fetchErr)
	}
	if len(messages) != 0 {
		t.Fatalf("fetchKafkaMessageBatchWith() messages = %#v, want none", messages)
	}
}

func TestProcessKafkaMessagePersistBatchFlushesClaimsBeforePrepareRetry(t *testing.T) {
	prepareErr := errors.New("redis temporarily unavailable")
	prepareCalls := 0
	persistCalls := 0

	err := processKafkaMessagePersistBatch(context.Background(), []kafka.Message{{Offset: 10}, {Offset: 11}}, "group_message_persist",
		func(_ context.Context, message kafka.Message, topic string) (*kafkaMessagePersistItem, bool, error) {
			prepareCalls++
			if topic != "group_message_persist" {
				t.Fatalf("prepare topic = %q", topic)
			}
			if message.Offset == 11 {
				return nil, false, prepareErr
			}
			return &kafkaMessagePersistItem{
				tenantID: types.TenantID(7),
				payload: kafkaMessagePersistPayload{
					MessageKind: "group",
					ShardMonth:  "202610",
					Message:     types.Message{ShardMonth: "202610"},
				},
			}, false, nil
		},
		func(_ context.Context, groups map[string][]*kafkaMessagePersistItem) error {
			persistCalls++
			items := groups["7:202610"]
			if len(items) != 1 {
				t.Fatalf("persisted groups = %#v, want one claimed item", groups)
			}
			return nil
		})

	if !errors.Is(err, prepareErr) {
		t.Fatalf("processKafkaMessagePersistBatch() error = %v, want %v", err, prepareErr)
	}
	if prepareCalls != 2 || persistCalls != 1 {
		t.Fatalf("calls prepare=%d persist=%d, want 2/1", prepareCalls, persistCalls)
	}
}

func TestProcessKafkaMessagePersistBatchReportsFlushAndPrepareErrors(t *testing.T) {
	prepareErr := errors.New("claim failed")
	persistErr := errors.New("database unavailable")
	prepareCalls := 0

	err := processKafkaMessagePersistBatch(context.Background(), []kafka.Message{{Offset: 10}, {Offset: 11}}, "message_persist",
		func(_ context.Context, _ kafka.Message, _ string) (*kafkaMessagePersistItem, bool, error) {
			prepareCalls++
			if prepareCalls == 2 {
				return nil, false, prepareErr
			}
			return &kafkaMessagePersistItem{tenantID: types.TenantID(7)}, false, nil
		},
		func(context.Context, map[string][]*kafkaMessagePersistItem) error {
			return persistErr
		})

	if !errors.Is(err, persistErr) || !strings.Contains(err.Error(), prepareErr.Error()) {
		t.Fatalf("processKafkaMessagePersistBatch() error = %v, want both flush and prepare errors", err)
	}
}

func TestConsumeKafkaMalformedMessageRetriesCommit(t *testing.T) {
	commitCalls := 0
	waitCalls := 0
	handlerCalls := 0
	err := consumeKafkaMessage(context.Background(), kafka.Message{Value: []byte("{")},
		func(context.Context, kafkaEnvelope, kafka.Message) error {
			handlerCalls++
			return nil
		},
		func(context.Context, kafka.Message) error {
			commitCalls++
			if commitCalls == 1 {
				return errors.New("commit failed")
			}
			return nil
		},
		func(context.Context, time.Duration) error {
			waitCalls++
			return nil
		})

	if err != nil {
		t.Fatalf("consumeKafkaMessage() error = %v", err)
	}
	if handlerCalls != 0 || commitCalls != 2 || waitCalls != 1 {
		t.Fatalf("calls handler=%d commit=%d wait=%d, want 0/2/1", handlerCalls, commitCalls, waitCalls)
	}
}

func TestRunKafkaMessageSingleRetriesReturnsErrorWithoutTerminalCommit(t *testing.T) {
	operationCalls := 0
	markFailedCalls := 0
	waitCalls := 0

	err := runKafkaMessageSingleRetries(context.Background(), 3,
		func() error {
			operationCalls++
			return errors.New("write failed")
		},
		func(error) (*types.KafkaConsumeRecord, error) {
			markFailedCalls++
			return &types.KafkaConsumeRecord{RetryCount: markFailedCalls}, nil
		},
		func(context.Context, time.Duration) error {
			waitCalls++
			return nil
		})

	if err == nil || !strings.Contains(err.Error(), "message persistence retries exhausted") {
		t.Fatalf("runKafkaMessageSingleRetries() error = %v, want retry exhaustion", err)
	}
	if operationCalls != 4 || markFailedCalls != 4 || waitCalls != 3 {
		t.Fatalf("calls operation=%d failed=%d wait=%d, want 4/4/3",
			operationCalls, markFailedCalls, waitCalls)
	}
}

func TestRunKafkaMessageSingleRetriesStopsOnSuccess(t *testing.T) {
	operationCalls := 0
	markFailedCalls := 0
	err := runKafkaMessageSingleRetries(context.Background(), 3,
		func() error {
			operationCalls++
			if operationCalls < 3 {
				return errors.New("temporary failure")
			}
			return nil
		},
		func(error) (*types.KafkaConsumeRecord, error) {
			markFailedCalls++
			return &types.KafkaConsumeRecord{RetryCount: markFailedCalls}, nil
		},
		func(context.Context, time.Duration) error { return nil })

	if err != nil {
		t.Fatalf("runKafkaMessageSingleRetries() error = %v", err)
	}
	if operationCalls != 3 || markFailedCalls != 2 {
		t.Fatalf("calls operation=%d failed=%d, want 3/2", operationCalls, markFailedCalls)
	}
}