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

pubsub/tests/bench_pubsub.go

@@ -0,0 +1,41 @@
+package tests
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/ThreeDotsLabs/watermill/message/subscriber"
+)
+
+// BenchmarkPubSubConstructor is a function that creates a Publisher and Subscriber to be used for benchmarks.
+type BenchmarkPubSubConstructor func(n int) (message.Publisher, message.Subscriber)
+
+// BenchSubscriber runs benchmark on a message Subscriber.
+func BenchSubscriber(b *testing.B, pubSubConstructor BenchmarkPubSubConstructor) {
+	pub, sub := pubSubConstructor(b.N)
+	topicName := testTopicName(TestContext{TestID: NewTestID()})
+
+	messages, err := sub.Subscribe(context.Background(), topicName)
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	go func() {
+		for i := 0; i < b.N; i++ {
+			msg := message.NewMessage("1", nil)
+			err := pub.Publish(topicName, msg)
+			if err != nil {
+				panic(err)
+			}
+		}
+	}()
+
+	b.ResetTimer()
+
+	consumedMessages, all := subscriber.BulkRead(messages, b.N, time.Second*60)
+	if !all {
+		b.Fatalf("not all messages received, have %d, expected %d", len(consumedMessages), b.N)
+	}
+}

pubsub/tests/test_asserts.go

@@ -0,0 +1,106 @@
+package tests
+
+import (
+	"context"
+	"sort"
+	"testing"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/stretchr/testify/assert"
+)
+
+func difference(a, b []string) []string {
+	mb := map[string]bool{}
+	for _, x := range b {
+		mb[x] = true
+	}
+	ab := []string{}
+	for _, x := range a {
+		if _, ok := mb[x]; !ok {
+			ab = append(ab, x)
+		}
+	}
+	return ab
+}
+
+// MissingMessages returns a list of missing messages UUIDs.
+func MissingMessages(expected message.Messages, received message.Messages) []string {
+	sentIDs := expected.IDs()
+	receivedIDs := received.IDs()
+
+	sort.Strings(sentIDs)
+	sort.Strings(receivedIDs)
+
+	return difference(sentIDs, receivedIDs)
+}
+
+// AssertAllMessagesReceived checks if all messages were received,
+// ignoring the order and assuming that they are already deduplicated.
+func AssertAllMessagesReceived(t *testing.T, sent message.Messages, received message.Messages) bool {
+	sentIDs := sent.IDs()
+	receivedIDs := received.IDs()
+
+	sort.Strings(sentIDs)
+	sort.Strings(receivedIDs)
+
+	if len(sentIDs) != len(receivedIDs) {
+		t.Errorf("id's count is different: received: %d, sent: %d", len(receivedIDs), len(sentIDs))
+	}
+
+	missing := MissingMessages(sent, received)
+	extra := MissingMessages(received, sent)
+
+	if len(missing) > 0 || len(extra) > 0 {
+		t.Errorf("received different messages ID's, missing: %s, extra %s", missing, extra)
+		return false
+	}
+
+	return true
+}
+
+// AssertMessagesPayloads check if received messages have the same payload as expected in expectedPayloads.
+func AssertMessagesPayloads(
+	t *testing.T,
+	expectedPayloads map[string][]byte,
+	received []*message.Message,
+) bool {
+	assert.Len(t, received, len(expectedPayloads))
+
+	receivedMsgs := map[string]interface{}{}
+	for _, msg := range received {
+		receivedMsgs[msg.UUID] = string(msg.Payload)
+	}
+
+	ok := true
+	for msgUUID, sentMsgPayload := range expectedPayloads {
+		if !assert.EqualValues(t, sentMsgPayload, receivedMsgs[msgUUID]) {
+			ok = false
+		}
+	}
+
+	return ok
+}
+
+// AssertMessagesMetadata checks if metadata of all received messages is the same as in expectedValues.
+func AssertMessagesMetadata(t *testing.T, key string, expectedValues map[string]string, received []*message.Message) bool {
+	assert.Len(t, received, len(expectedValues))
+
+	ok := true
+	for _, msg := range received {
+		if !assert.Equal(t, expectedValues[msg.UUID], msg.Metadata[key]) {
+			ok = false
+		}
+	}
+
+	return ok
+}
+
+// AssertAllMessagesHaveSameContext checks if context of all received messages is the same as in expectedValues, if PreserveContext is enabled.
+func AssertAllMessagesHaveSameContext(t *testing.T, contextKeyString string, expectedValues map[string]context.Context, received []*message.Message) {
+	assert.Len(t, received, len(expectedValues))
+	for _, msg := range received {
+		expectedValue := expectedValues[msg.UUID].Value(contextKey(contextKeyString)).(string)
+		actualValue := msg.Context().Value(contextKey(contextKeyString))
+		assert.Equal(t, expectedValue, actualValue)
+	}
+}

pubsub/tests/test_pubsub_stress.go

@@ -0,0 +1,16 @@
+//go:build stress
+// +build stress
+
+package tests
+
+import (
+	"runtime"
+)
+
+func init() {
+	// stress tests may work a bit slower
+	defaultTimeout *= 5
+
+	// Set GOMAXPROCS to double the number of CPUs
+	runtime.GOMAXPROCS(runtime.GOMAXPROCS(0) * 2)
+}