upstream-watermill/message/subscriber/read_test.go

package subscriber_test

import (
	"testing"
	"time"

	"github.com/ThreeDotsLabs/watermill/pubsub/tests"

	"github.com/ThreeDotsLabs/watermill"
	"github.com/ThreeDotsLabs/watermill/message"
	"github.com/ThreeDotsLabs/watermill/message/subscriber"
	"github.com/stretchr/testify/assert"
)

type bulkReadFunc func(messagesCh <-chan *message.Message, limit int, timeout time.Duration) (receivedMessages message.Messages, all bool)

func TestBulkRead(t *testing.T) {
	testCases := []struct {
		Name         string
		BulkReadFunc bulkReadFunc
	}{
		{
			Name:         "BulkRead",
			BulkReadFunc: subscriber.BulkRead,
		},
		{
			Name:         "BulkReadWithDeduplication",
			BulkReadFunc: subscriber.BulkReadWithDeduplication,
		},
	}

	for _, c := range testCases {
		t.Run(c.Name, func(t *testing.T) {
			messagesCount := 100

			var messages []*message.Message
			messagesCh := make(chan *message.Message, messagesCount)

			for i := 0; i < messagesCount; i++ {
				msg := message.NewMessage(watermill.NewUUID(), nil)

				messages = append(messages, msg)
				messagesCh <- msg
			}

			readMessages, all := subscriber.BulkRead(messagesCh, messagesCount, time.Second)
			assert.True(t, all)

			tests.AssertAllMessagesReceived(t, messages, readMessages)
		})
	}
}

func TestBulkRead_timeout(t *testing.T) {
	testCases := []struct {
		Name         string
		BulkReadFunc bulkReadFunc
	}{
		{
			Name:         "BulkRead",
			BulkReadFunc: subscriber.BulkRead,
		},
		{
			Name:         "BulkReadWithDeduplication",
			BulkReadFunc: subscriber.BulkReadWithDeduplication,
		},
	}

	for _, c := range testCases {
		t.Run(c.Name, func(t *testing.T) {
			messagesCount := 100
			sendLimit := 90

			messagesCh := make(chan *message.Message, messagesCount)

			for i := 0; i < messagesCount; i++ {
				msg := message.NewMessage(watermill.NewUUID(), nil)

				if i < sendLimit {
					messagesCh <- msg
				}
			}

			bulkReadStart := time.Now()
			readMessages, all := subscriber.BulkRead(messagesCh, messagesCount, time.Millisecond)

			assert.WithinDuration(t, bulkReadStart, time.Now(), time.Millisecond*100)
			assert.False(t, all)
			assert.Equal(t, sendLimit, len(readMessages))
		})
	}
}

func TestBulkRead_with_limit(t *testing.T) {
	testCases := []struct {
		Name         string
		BulkReadFunc bulkReadFunc
	}{
		{
			Name:         "BulkRead",
			BulkReadFunc: subscriber.BulkRead,
		},
		{
			Name:         "BulkReadWithDeduplication",
			BulkReadFunc: subscriber.BulkReadWithDeduplication,
		},
	}

	for _, c := range testCases {
		t.Run(c.Name, func(t *testing.T) {
			messagesCount := 110
			limit := 100

			messagesCh := make(chan *message.Message, messagesCount)

			for i := 0; i < messagesCount; i++ {
				msg := message.NewMessage(watermill.NewUUID(), nil)

				messagesCh <- msg
			}

			readMessages, all := subscriber.BulkRead(messagesCh, limit, time.Second)
			assert.True(t, all)
			assert.Equal(t, limit, len(readMessages))
		})
	}
}

func TestBulkRead_return_on_channel_close(t *testing.T) {
	testCases := []struct {
		Name         string
		BulkReadFunc bulkReadFunc
	}{
		{
			Name:         "BulkRead",
			BulkReadFunc: subscriber.BulkRead,
		},
		{
			Name:         "BulkReadWithDeduplication",
			BulkReadFunc: subscriber.BulkReadWithDeduplication,
		},
	}

	for _, c := range testCases {
		t.Run(c.Name, func(t *testing.T) {
			messagesCount := 100
			sendLimit := 90

			messagesCh := make(chan *message.Message, messagesCount)
			messagesChClosed := false

			for i := 0; i < messagesCount; i++ {
				msg := message.NewMessage(watermill.NewUUID(), nil)

				if i < sendLimit {
					messagesCh <- msg
				} else if !messagesChClosed {
					close(messagesCh)
					messagesChClosed = true
				}
			}

			bulkReadStart := time.Now()
			_, all := subscriber.BulkRead(messagesCh, messagesCount, time.Second)

			assert.WithinDuration(t, bulkReadStart, time.Now(), time.Millisecond*100)
			assert.False(t, all)
		})
	}
}

func TestBulkReadWithDeduplication(t *testing.T) {
	messagesCh := make(chan *message.Message, 3)

	msg1 := message.NewMessage(watermill.NewUUID(), nil)
	msg2 := message.NewMessage(watermill.NewUUID(), nil)
	messagesCh <- msg1
	messagesCh <- msg1
	messagesCh <- msg2

	readMessages, all := subscriber.BulkReadWithDeduplication(messagesCh, 2, time.Second)
	assert.True(t, all)

	assert.Equal(t, []string{msg1.UUID, msg2.UUID}, readMessages.IDs())
}