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 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)
	}
}