snapshot of https://github.com/ThreeDotsLabs/watermill.git (master, 2026-09-15; blobs > 131072B omitted)

pubsub/gochannel/pubsub_test.go

@@ -0,0 +1,399 @@
+package gochannel_test
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"sync"
+	"testing"
+	"time"
+
+	"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"
+	"github.com/ThreeDotsLabs/watermill/pubsub/tests"
+)
+
+func createPersistentPubSub(t *testing.T) (message.Publisher, message.Subscriber) {
+	pubSub := gochannel.NewGoChannel(
+		gochannel.Config{
+			OutputChannelBuffer: 10000,
+			Persistent:          true,
+		},
+		watermill.NewStdLogger(true, true),
+	)
+	return pubSub, pubSub
+}
+
+func createPersistentPubSubWithContextPreserved(t *testing.T) (message.Publisher, message.Subscriber) {
+	pubSub := gochannel.NewGoChannel(
+		gochannel.Config{
+			OutputChannelBuffer: 10000,
+			Persistent:          true,
+			PreserveContext:     true,
+		},
+		watermill.NewStdLogger(true, true),
+	)
+	return pubSub, pubSub
+}
+
+func TestPublishSubscribe_persistent(t *testing.T) {
+	tests.TestPubSub(
+		t,
+		tests.Features{
+			ConsumerGroups:        false,
+			ExactlyOnceDelivery:   true,
+			GuaranteedOrder:       false,
+			Persistent:            false,
+			RequireSingleInstance: true,
+		},
+		createPersistentPubSub,
+		nil,
+	)
+}
+
+func TestPublishSubscribe_context_preserved(t *testing.T) {
+	tests.TestPubSub(
+		t,
+		tests.Features{
+			ConsumerGroups:        false,
+			ExactlyOnceDelivery:   true,
+			GuaranteedOrder:       false,
+			Persistent:            false,
+			RequireSingleInstance: true,
+			ContextPreserved:      true,
+		},
+		createPersistentPubSubWithContextPreserved,
+		nil,
+	)
+}
+
+func TestPublishSubscribe_not_persistent(t *testing.T) {
+	messagesCount := 100
+	pubSub := gochannel.NewGoChannel(
+		gochannel.Config{OutputChannelBuffer: int64(messagesCount)},
+		watermill.NewStdLogger(true, true),
+	)
+	topicName := "test_topic_" + watermill.NewUUID()
+
+	msgs, err := pubSub.Subscribe(context.Background(), topicName)
+	require.NoError(t, err)
+
+	sendMessages := tests.PublishSimpleMessages(t, messagesCount, pubSub, topicName)
+	receivedMsgs, _ := subscriber.BulkRead(msgs, messagesCount, time.Second)
+
+	tests.AssertAllMessagesReceived(t, sendMessages, receivedMsgs)
+
+	assert.NoError(t, pubSub.Close())
+}
+
+func TestPublishSubscribe_not_persistent_with_context(t *testing.T) {
+	messagesCount := 100
+	pubSub := gochannel.NewGoChannel(
+		gochannel.Config{OutputChannelBuffer: int64(messagesCount), PreserveContext: true},
+		watermill.NewStdLogger(true, true),
+	)
+	topicName := "test_topic_" + watermill.NewUUID()
+
+	msgs, err := pubSub.Subscribe(context.Background(), topicName)
+	require.NoError(t, err)
+
+	const contextKeyString = "foo"
+	sendMessages := tests.PublishSimpleMessagesWithContext(t, messagesCount, contextKeyString, pubSub, topicName)
+	receivedMsgs, _ := subscriber.BulkRead(msgs, messagesCount, time.Second)
+
+	expectedContexts := make(map[string]context.Context)
+	for _, msg := range sendMessages {
+		expectedContexts[msg.UUID] = msg.Context()
+	}
+	tests.AssertAllMessagesReceived(t, sendMessages, receivedMsgs)
+	tests.AssertAllMessagesHaveSameContext(t, contextKeyString, expectedContexts, receivedMsgs)
+
+	assert.NoError(t, pubSub.Close())
+}
+
+func TestPublishSubscribe_block_until_ack(t *testing.T) {
+	pubSub := gochannel.NewGoChannel(
+		gochannel.Config{BlockPublishUntilSubscriberAck: true},
+		watermill.NewStdLogger(true, true),
+	)
+	topicName := "test_topic_" + watermill.NewUUID()
+
+	msgs, err := pubSub.Subscribe(context.Background(), topicName)
+	require.NoError(t, err)
+
+	published := make(chan struct{})
+	go func() {
+		err := pubSub.Publish(topicName, message.NewMessage("1", nil))
+		require.NoError(t, err)
+		close(published)
+	}()
+
+	msg1 := <-msgs
+	select {
+	case <-published:
+		t.Fatal("publish should be blocked until ack")
+	default:
+		// ok
+	}
+
+	msg1.Nack()
+	select {
+	case <-published:
+		t.Fatal("publish should be blocked after nack")
+	default:
+		// ok
+	}
+
+	msg2 := <-msgs
+	msg2.Ack()
+
+	select {
+	case <-published:
+		// ok
+	case <-time.After(time.Second):
+		t.Fatal("publish should be not blocked after ack")
+	}
+}
+
+func TestPublishSubscribe_race_condition_on_subscribe(t *testing.T) {
+	testsCount := 15
+	if testing.Short() {
+		testsCount = 3
+	}
+
+	for i := 0; i < testsCount; i++ {
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			t.Parallel()
+			testPublishSubscribeSubRace(t)
+		})
+	}
+}
+
+func TestSubscribe_race_condition_when_closing(t *testing.T) {
+	testsCount := 15
+	if testing.Short() {
+		testsCount = 3
+	}
+
+	for i := 0; i < testsCount; i++ {
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			t.Parallel()
+			pubSub := gochannel.NewGoChannel(
+				gochannel.Config{},
+				watermill.NewStdLogger(true, false),
+			)
+			go func() {
+				err := pubSub.Close()
+				require.NoError(t, err)
+			}()
+			_, err := pubSub.Subscribe(context.Background(), "topic")
+			require.NoError(t, err)
+		})
+	}
+}
+
+func TestPublish_race_condition_when_closing(t *testing.T) {
+	testsCount := 15
+	if testing.Short() {
+		testsCount = 3
+	}
+
+	for i := 0; i < testsCount; i++ {
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			t.Parallel()
+			pubSub := gochannel.NewGoChannel(
+				gochannel.Config{},
+				watermill.NewStdLogger(true, false),
+			)
+			go func() {
+				_ = pubSub.Publish("topic", message.NewMessage(watermill.NewShortUUID(), nil))
+			}()
+
+			err := pubSub.Close()
+			require.NoError(t, err)
+		})
+	}
+}
+
+func TestPublishSubscribe_do_not_block_other_subscribers(t *testing.T) {
+	pubSub := gochannel.NewGoChannel(
+		gochannel.Config{},
+		watermill.NewStdLogger(true, true),
+	)
+	topicName := "test_topic_" + watermill.NewUUID()
+
+	msgsFromSubscriber1, err := pubSub.Subscribe(context.Background(), topicName)
+	require.NoError(t, err)
+
+	_, err = pubSub.Subscribe(context.Background(), topicName)
+	require.NoError(t, err)
+
+	msgsFromSubscriber3, err := pubSub.Subscribe(context.Background(), topicName)
+	require.NoError(t, err)
+
+	err = pubSub.Publish(topicName, message.NewMessage("1", nil))
+	require.NoError(t, err)
+
+	received := make(chan struct{})
+	go func() {
+		msg := <-msgsFromSubscriber1
+		msg.Ack()
+
+		msg = <-msgsFromSubscriber3
+		msg.Ack()
+
+		close(received)
+	}()
+
+	select {
+	case <-received:
+		// ok
+	case <-time.After(5 * time.Second):
+		t.Fatal("subscriber which didn't ack a message blocked other subscribers from receiving it")
+	}
+}
+
+func TestPublishSubscribe_flush_output_channel(t *testing.T) {
+	messagesCount := 300
+	logger := watermill.NewStdLogger(true, true)
+	ctx := context.Background()
+	config := gochannel.Config{
+		OutputChannelBuffer:            int64(messagesCount),
+		Persistent:                     false,
+		BlockPublishUntilSubscriberAck: false,
+	}
+	pubSub := gochannel.NewGoChannel(
+		config,
+		logger,
+	)
+
+	totalMessage := 0
+	artificialWorkload := time.Millisecond * 5 //keep it small but noticeable in logs
+	topicName := "test_topic"
+
+	messageChannel, err := pubSub.Subscribe(ctx, topicName)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// waitgroup for stopping the subscriber handler from processing/receiving messages until we are done filling the buffer of the pubSub
+	var wgStartSubscriber sync.WaitGroup
+	wgStartSubscriber.Add(1)
+	// waitgroup for expected buffer to be flushed
+	var wgFlushBuffer sync.WaitGroup
+	wgFlushBuffer.Add(messagesCount)
+
+	// start subscriber handler in a go routine
+	// reads out the messages, if all is ok it should be able to ("flush") read all messages from a 'closed' pubsub
+	go func(messageChannel <-chan *message.Message) {
+		wgStartSubscriber.Wait()
+		for msg := range messageChannel {
+			// artificial workload
+			time.Sleep(artificialWorkload)
+			msg.Ack()
+			logger.Trace("message acked", nil)
+			// would normally use atomic value here but concurrency shouldn't be an issue for this test
+			totalMessage++
+			wgFlushBuffer.Done()
+		}
+		logger.Trace("channel closed", nil)
+	}(messageChannel)
+
+	tests.PublishSimpleMessages(t, messagesCount, pubSub, topicName)
+	// wait for buffer to fill then start reading in subscriber handler
+	for {
+		if len(messageChannel) != int(config.OutputChannelBuffer) {
+			continue
+		} else {
+			wgStartSubscriber.Done()
+			break
+		}
+	}
+
+	err = pubSub.Close()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Publishing new message should still error as expected
+	err = pubSub.Publish(topicName, message.NewMessage(watermill.NewUUID(), []byte("x")))
+	assert.ErrorContains(t, err, "Pub/Sub closed")
+
+	// And so should subscribe
+	_, err = pubSub.Subscribe(ctx, topicName)
+	assert.ErrorContains(t, err, "Pub/Sub closed")
+
+	wgFlushBuffer.Wait()
+	// But subscriber handler should still be able to read the remaining messages from output channel aka flushing
+	assert.Equal(t, messagesCount, totalMessage)
+}
+
+func testPublishSubscribeSubRace(t *testing.T) {
+	t.Helper()
+
+	messagesCount := 500
+	subscribersCount := 200
+	if testing.Short() {
+		messagesCount = 200
+		subscribersCount = 20
+	}
+
+	pubSub := gochannel.NewGoChannel(
+		gochannel.Config{
+			OutputChannelBuffer: int64(messagesCount),
+			Persistent:          true,
+		},
+		watermill.NewStdLogger(true, false),
+	)
+
+	allSent := sync.WaitGroup{}
+	allSent.Add(messagesCount)
+	allReceived := sync.WaitGroup{}
+
+	sentMessages := message.Messages{}
+	go func() {
+		for i := 0; i < messagesCount; i++ {
+			msg := message.NewMessage(watermill.NewUUID(), nil)
+			sentMessages = append(sentMessages, msg)
+
+			go func() {
+				require.NoError(t, pubSub.Publish("topic", msg))
+				allSent.Done()
+			}()
+		}
+	}()
+
+	subscriberReceivedCh := make(chan message.Messages, subscribersCount)
+	for i := 0; i < subscribersCount; i++ {
+		allReceived.Add(1)
+
+		go func() {
+			msgs, err := pubSub.Subscribe(context.Background(), "topic")
+			require.NoError(t, err)
+
+			received, _ := subscriber.BulkRead(msgs, messagesCount, time.Second*10)
+			subscriberReceivedCh <- received
+
+			allReceived.Done()
+		}()
+	}
+
+	log.Println("waiting for all sent")
+	allSent.Wait()
+
+	log.Println("waiting for all received")
+	allReceived.Wait()
+
+	close(subscriberReceivedCh)
+
+	log.Println("asserting")
+
+	for subMsgs := range subscriberReceivedCh {
+		tests.AssertAllMessagesReceived(t, sentMessages, subMsgs)
+	}
+}