Add Test cases

message/pubsub.go

@@ -31,7 +31,7 @@ type Subscriber interface {
 	// If message processing fails and the message should be redelivered `Nack()` should be called instead.
 	//
 	// When the provided ctx is canceled, the subscriber closes the subscription and the output channel.
-	// The provided ctx is passed to all produced messages.
+	// The provided ctx is passed to all produced messages (this is configurable for the local Pub/Sub implementation).
 	// When Nack or Ack is called on the message, the context of the message is canceled.
 	Subscribe(ctx context.Context, topic string) (<-chan *Message, error)
 	// Close closes all subscriptions with their output channels and flushes offsets etc. when needed.

pubsub/gochannel/pubsub_test.go

@@ -29,6 +29,18 @@ func createPersistentPubSub(t *testing.T) (message.Publisher, message.Subscriber
 	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,
@@ -44,6 +56,22 @@ func TestPublishSubscribe_persistent(t *testing.T) {
 	)
 }
 
+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(
@@ -63,6 +91,31 @@ func TestPublishSubscribe_not_persistent(t *testing.T) {
 	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 contextKey = "foo"
+	sendMessages := tests.PublishSimpleMessagesWithContext(t, messagesCount, contextKey, 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, contextKey, expectedContexts, receivedMsgs)
+
+	assert.NoError(t, pubSub.Close())
+}
+
 func TestPublishSubscribe_block_until_ack(t *testing.T) {
 	pubSub := gochannel.NewGoChannel(
 		gochannel.Config{BlockPublishUntilSubscriberAck: true},

pubsub/tests/test_asserts.go

@@ -1,6 +1,7 @@
 package tests
 
 import (
+	"context"
 	"sort"
 	"testing"
 
@@ -92,3 +93,19 @@ func AssertMessagesMetadata(t *testing.T, key string, expectedValues map[string]
 
 	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, contextKey string, expectedValues map[string]context.Context, received []*message.Message) bool {
+	assert.Len(t, received, len(expectedValues))
+
+	ok := true
+	for _, msg := range received {
+		expectedValue := expectedValues[msg.UUID].Value(contextKey)
+		actualValue := msg.Context().Value(contextKey)
+		if !assert.Equal(t, expectedValue, actualValue) {
+			ok = false
+		}
+	}
+
+	return ok
+}

pubsub/tests/test_pubsub.go

@@ -124,6 +124,10 @@ type Features struct {
 
 	// GenerateTopicFunc overrides standard topic name generation.
 	GenerateTopicFunc func(tctx TestContext) string
+
+	// ContextPreserved should be set to true if the Pub/Sub implementation preserves the context
+	// of the message when it's published and consumed.
+	ContextPreserved bool
 }
 
 // RunOnlyFastTests returns true if -short flag was provided -race was not provided.
@@ -993,7 +997,14 @@ func TestSubscribeCtx(
 	if subscribeInitializer, ok := sub.(message.SubscribeInitializer); ok {
 		require.NoError(t, subscribeInitializer.SubscribeInitialize(topicName))
 	}
-	publishedMessages := PublishSimpleMessages(t, messagesCount, pub, topicName)
+
+	var publishedMessages message.Messages
+	var contextKeyString = "abc"
+	if tCtx.Features.ContextPreserved {
+		publishedMessages = PublishSimpleMessagesWithContext(t, messagesCount, contextKeyString, pub, topicName)
+	} else {
+		publishedMessages = PublishSimpleMessages(t, messagesCount, pub, topicName)
+	}
 
 	msgsToCancel, err := sub.Subscribe(ctxWithCancel, topicName)
 	require.NoError(t, err)
@@ -1017,15 +1028,23 @@ ClosedLoop:
 	}
 
 	ctx := context.WithValue(context.Background(), contextKey("foo"), "bar")
+
+	// For mocking the output of pub-subs where context is preserved vs not preserved
+	expectedContexts := make(map[string]context.Context)
+	for _, msg := range publishedMessages {
+		if tCtx.Features.ContextPreserved {
+			expectedContexts[msg.UUID] = msg.Context()
+		} else {
+			expectedContexts[msg.UUID] = ctx
+		}
+	}
+
 	msgs, err := sub.Subscribe(ctx, topicName)
 	require.NoError(t, err)
 
 	receivedMessages, _ := bulkRead(tCtx, msgs, messagesCount, defaultTimeout)
 	AssertAllMessagesReceived(t, publishedMessages, receivedMessages)
-
-	for _, msg := range receivedMessages {
-		assert.EqualValues(t, "bar", msg.Context().Value(contextKey("foo")))
-	}
+	AssertAllMessagesHaveSameContext(t, contextKeyString, expectedContexts, receivedMessages)
 }
 
 // TestReconnect tests if reconnecting to a Pub/Sub works correctly.
@@ -1271,6 +1290,24 @@ func PublishSimpleMessages(t *testing.T, messagesCount int, publisher message.Pu
 	return messagesToPublish
 }
 
+// PublishSimpleMessagesWithContext publishes provided number of simple messages without a payload, but custom context
+func PublishSimpleMessagesWithContext(t *testing.T, messagesCount int, contextKey string, publisher message.Publisher, topicName string) message.Messages {
+	var messagesToPublish []*message.Message
+
+	for i := 0; i < messagesCount; i++ {
+		id := watermill.NewUUID()
+
+		msg := message.NewMessage(id, nil)
+		msg.SetContext(context.WithValue(context.Background(), contextKey, "bar"+strconv.Itoa(i)))
+		messagesToPublish = append(messagesToPublish, msg)
+
+		err := publishWithRetry(publisher, topicName, msg)
+		require.NoError(t, err, "cannot publish messages")
+	}
+
+	return messagesToPublish
+}
+
 // AddSimpleMessagesParallel publishes provided number of simple messages without a payload
 // using the provided number of publishers (goroutines).
 func AddSimpleMessagesParallel(t *testing.T, messagesCount int, publisher message.Publisher, topicName string, publishers int) message.Messages {