101 lines
2.2 KiB
Go
101 lines
2.2 KiB
Go
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()
|
|
}
|