upstream-watermill/message/decorator_test.go

package message_test

import (
	"context"
	"strconv"
	"testing"
	"time"

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

	"github.com/pkg/errors"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/ThreeDotsLabs/watermill"
	"github.com/ThreeDotsLabs/watermill/message"
	"github.com/ThreeDotsLabs/watermill/message/subscriber"
	"github.com/ThreeDotsLabs/watermill/pubsub/gochannel"
)

var noop = func(*message.Message) {}
var closingErr = errors.New("mock error on close")

type mockSubscriber struct {
	ch chan *message.Message
}

func (m mockSubscriber) Subscribe(context.Context, string) (<-chan *message.Message, error) {
	return m.ch, nil
}

func (m mockSubscriber) Close() error {
	close(m.ch)
	return nil
}

func TestMessageTransformSubscriberDecorator_transparent(t *testing.T) {
	sub := mockSubscriber{make(chan *message.Message)}
	decorated, err := message.MessageTransformSubscriberDecorator(noop)(sub)
	require.NoError(t, err)

	messages, err := decorated.Subscribe(context.Background(), "topic")
	require.NoError(t, err)

	richMessage := message.NewMessage("uuid", []byte("serious payloads"))
	richMessage.Metadata.Set("k1", "v1")
	richMessage.Metadata.Set("k2", "v2")

	go func() {
		sub.ch <- richMessage
	}()

	received, all := subscriber.BulkRead(messages, 1, time.Second)
	require.True(t, all)

	assert.True(t, received[0].Equals(richMessage), "expected the message to pass unchanged through decorator")
}

type closingSubscriber struct {
	closed bool
}

func (closingSubscriber) Subscribe(context.Context, string) (<-chan *message.Message, error) {
	return nil, nil
}

func (c *closingSubscriber) Close() error {
	c.closed = true
	return closingErr
}

func TestMessageTransformSubscriberDecorator_Close(t *testing.T) {
	cs := &closingSubscriber{}

	decoratedSub, err := message.MessageTransformSubscriberDecorator(noop)(cs)
	require.NoError(t, err)

	// given
	require.False(t, cs.closed)

	// when
	decoratedCloseErr := decoratedSub.Close()

	// then
	assert.True(
		t,
		cs.closed,
		"expected the Close() call to propagate to decorated subscriber",
	)
	assert.Equal(
		t,
		closingErr,
		decoratedCloseErr,
		"expected the decorator to propagate the closing error from underlying subscriber",
	)
}

func TestMessageTransformSubscriberDecorator_Subscribe(t *testing.T) {
	numMessages := 1000
	pubSub := gochannel.NewGoChannel(gochannel.Config{}, watermill.NewStdLogger(true, true))

	onMessage := func(msg *message.Message) {
		msg.Metadata.Set("key", "value")
	}
	decorator := message.MessageTransformSubscriberDecorator(onMessage)

	decoratedSub, err := decorator(pubSub)
	require.NoError(t, err)

	messages, err := decoratedSub.Subscribe(context.Background(), "topic")
	require.NoError(t, err)

	sent := message.Messages{}

	go func() {
		for i := 0; i < numMessages; i++ {
			msg := message.NewMessage(strconv.Itoa(i), []byte{})
			sent = append(sent, msg)

			err = pubSub.Publish("topic", msg)
			require.NoError(t, err)
		}
	}()

	received, all := subscriber.BulkRead(messages, numMessages, time.Second)
	require.True(t, all)
	tests.AssertAllMessagesReceived(t, sent, received)

	for _, msg := range received {
		assert.Equal(
			t,
			"value",
			msg.Metadata.Get("key"),
			"expected onMessage callback to have set metadata",
		)
	}
}

type mockPublisher struct {
	published message.Messages
}

func (m *mockPublisher) Publish(topic string, messages ...*message.Message) error {
	m.published = append(m.published, messages...)
	return nil
}

func (m mockPublisher) Close() error {
	return nil
}

func TestMessageTransformPublisherDecorator_transparent(t *testing.T) {
	pub := &mockPublisher{message.Messages{}}
	decorated, err := message.MessageTransformPublisherDecorator(noop)(pub)
	require.NoError(t, err)

	richMessage := message.NewMessage("uuid", []byte("serious payloads"))
	richMessage.Metadata.Set("k1", "v1")
	richMessage.Metadata.Set("k2", "v2")

	require.NoError(t, decorated.Publish("topic", richMessage))

	require.Len(t, pub.published, 1)
	assert.True(t, pub.published[0].Equals(richMessage), "expected the message to pass unchanged through decorator")
}

type closingPublisher struct {
	closed bool
}

func (c *closingPublisher) Publish(topic string, messages ...*message.Message) error {
	return nil
}

func (c *closingPublisher) Close() error {
	c.closed = true
	return closingErr
}

func TestMessageTransformPublisherDecorator_Close(t *testing.T) {
	cp := &closingPublisher{}

	decoratedPub, err := message.MessageTransformPublisherDecorator(noop)(cp)
	require.NoError(t, err)

	// given
	require.False(t, cp.closed)

	// when
	decoratedCloseErr := decoratedPub.Close()

	// then
	assert.True(
		t,
		cp.closed,
		"expected the Close() call to propagate to decorated publisher",
	)
	assert.Equal(
		t,
		closingErr,
		decoratedCloseErr,
		"expected the decorator to propagate the closing error from underlying publisher",
	)
}

func TestMessageTransformPublisherDecorator_Subscribe(t *testing.T) {
	numMessages := 1000
	pub := &mockPublisher{}

	onMessage := func(msg *message.Message) {
		msg.Metadata.Set("key", "value")
	}
	decorator := message.MessageTransformPublisherDecorator(onMessage)
	decoratedPub, err := decorator(pub)
	require.NoError(t, err)

	for i := 0; i < numMessages; i++ {
		msg := message.NewMessage(strconv.Itoa(i), []byte{})
		require.NoError(t, decoratedPub.Publish("topic", msg))
	}

	for i, msg := range pub.published {
		assert.Equal(t, strconv.Itoa(i), msg.UUID, "expected messages to arrive in unchanged order")
		assert.Equal(
			t,
			"value",
			msg.Metadata.Get("key"),
			"expected onMessage callback to have set metadata",
		)
	}
}

func TestMessageTransformer_nil_panics(t *testing.T) {
	require.Panics(
		t,
		func() {
			_ = message.MessageTransformSubscriberDecorator(nil)
		},
		"expected to panic if transform is nil",
	)
	require.Panics(
		t,
		func() {
			_ = message.MessageTransformPublisherDecorator(nil)
		},
		"expected to panic if transform is nil",
	)
}