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

pubsub/gochannel/pubsub_internal_test.go

@@ -0,0 +1,100 @@
+package gochannel
+
+import (
+	"context"
+	"strconv"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/ThreeDotsLabs/watermill"
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+func TestSubscribe_clean_subscriber_data(t *testing.T) {
+	subCount := 100
+	pubSub := NewGoChannel(
+		Config{OutputChannelBuffer: int64(subCount)},
+		watermill.NewStdLogger(false, false),
+	)
+	topicName := "test_topic"
+
+	for i := 0; i < subCount; i++ {
+		ctx, cancel := context.WithCancel(context.Background())
+		_, err := pubSub.Subscribe(ctx, topicName+"_index_"+strconv.Itoa(i))
+		require.NoError(t, err)
+		cancel()
+	}
+
+	err := pubSub.Close()
+	require.NoError(t, err)
+
+	assert.Len(t, pubSub.subscribers, 0)
+	lockCount := 0
+	pubSub.subscribersByTopicLock.Range(func(_, _ any) bool {
+		lockCount++
+		return true
+	})
+	assert.Equal(t, 0, lockCount)
+
+	assert.NoError(t, pubSub.Close())
+}
+
+func TestPublish_clean_lock_data(t *testing.T) {
+	messageCount := 100
+	pubSub := NewGoChannel(
+		Config{OutputChannelBuffer: int64(messageCount)},
+		watermill.NewStdLogger(false, false),
+	)
+	topicName := "test_topic"
+
+	_, err := pubSub.Subscribe(context.Background(), topicName+"_index_"+strconv.Itoa(0))
+	require.NoError(t, err)
+
+	for i := 0; i < messageCount; i++ {
+		err := pubSub.Publish(topicName+"_index_"+strconv.Itoa(i), message.NewMessage(watermill.NewShortUUID(), nil))
+		require.NoError(t, err)
+	}
+
+	lockCount := 0
+	pubSub.subscribersByTopicLock.Range(func(_, _ any) bool {
+		lockCount++
+		return true
+	})
+	assert.Equal(t, 1, lockCount)
+
+	assert.NoError(t, pubSub.Close())
+}
+
+// TestPublish_concurrent_close ensures Close does not race with in-flight persistent Publish calls. 
+func TestPublish_concurrent_close(t *testing.T) {
+	pubSub := NewGoChannel(
+		Config{Persistent: true},
+		watermill.NewStdLogger(false, false),
+	)
+	topicName := "test_topic"
+
+	publisherCount := 20
+	startPublishing := make(chan struct{})
+	wg := sync.WaitGroup{}
+
+	for i := 0; i < publisherCount; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			<-startPublishing
+			for {
+				err := pubSub.Publish(topicName, message.NewMessage(watermill.NewShortUUID(), nil))
+				if err != nil {
+					return
+				}
+			}
+		}()
+	}
+
+	close(startPublishing)
+	require.NoError(t, pubSub.Close())
+	wg.Wait()
+}