im_server/server/kafka_runtime_batch_test.go

package main

import (
	"context"
	"errors"
	"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 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 TestRunKafkaMessageSingleRetriesThenDeadLetters(t *testing.T) {
	operationCalls := 0
	markFailedCalls := 0
	markDeadCalls := 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(record *types.KafkaConsumeRecord, err error) error {
			markDeadCalls++
			if record == nil || record.RetryCount != 4 || err == nil {
				t.Fatalf("unexpected dead-letter input: record=%+v err=%v", record, err)
			}
			return nil
		},
		func(context.Context, time.Duration) error {
			waitCalls++
			return nil
		})

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

func TestRunKafkaMessageSingleRetriesStopsOnSuccess(t *testing.T) {
	operationCalls := 0
	markFailedCalls := 0
	markDeadCalls := 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(*types.KafkaConsumeRecord, error) error {
			markDeadCalls++
			return nil
		},
		func(context.Context, time.Duration) error { return nil })

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