package delay_test import ( "context" "fmt" "testing" "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/ThreeDotsLabs/watermill/components/delay" "github.com/ThreeDotsLabs/watermill/message" "github.com/ThreeDotsLabs/watermill/pubsub/gochannel" ) func TestPublisher(t *testing.T) { pubSub := gochannel.NewGoChannel(gochannel.Config{}, nil) messages, err := pubSub.Subscribe(context.Background(), "test") require.NoError(t, err) pub, err := delay.NewPublisher(pubSub, delay.PublisherConfig{}) require.NoError(t, err) pubAllowNoDelay, err := delay.NewPublisher(pubSub, delay.PublisherConfig{ AllowNoDelay: true, }) require.NoError(t, err) defaultDelayPub, err := delay.NewPublisher(pubSub, delay.PublisherConfig{ DefaultDelayGenerator: func(params delay.DefaultDelayGeneratorParams) (delay.Delay, error) { return delay.For(1 * time.Second), nil }, }) require.NoError(t, err) testCases := []struct { name string publisher message.Publisher messageConstructor func(id string) *message.Message expectedError bool expectedDelay time.Duration }{ { name: "no delay", publisher: pub, messageConstructor: func(id string) *message.Message { return message.NewMessage(id, nil) }, expectedError: true, expectedDelay: 0, }, { name: "no delay but allowed", publisher: pubAllowNoDelay, messageConstructor: func(id string) *message.Message { return message.NewMessage(id, nil) }, expectedDelay: 0, }, { name: "default delay", publisher: defaultDelayPub, messageConstructor: func(id string) *message.Message { return message.NewMessage(id, nil) }, expectedDelay: 1 * time.Second, }, { name: "delay from metadata", publisher: pub, messageConstructor: func(id string) *message.Message { msg := message.NewMessage(id, nil) delay.Message(msg, delay.For(2*time.Second)) return msg }, expectedDelay: 2 * time.Second, }, { name: "default delay override with metadata", publisher: defaultDelayPub, messageConstructor: func(id string) *message.Message { msg := message.NewMessage(id, nil) delay.Message(msg, delay.For(2*time.Second)) return msg }, expectedDelay: 2 * time.Second, }, { name: "delay from context", publisher: pub, messageConstructor: func(id string) *message.Message { msg := message.NewMessage(id, nil) ctx := delay.WithContext(context.Background(), delay.For(3*time.Second)) msg.SetContext(ctx) return msg }, expectedDelay: 3 * time.Second, }, { name: "default delay override with context", publisher: defaultDelayPub, messageConstructor: func(id string) *message.Message { msg := message.NewMessage(id, nil) ctx := delay.WithContext(context.Background(), delay.For(3*time.Second)) msg.SetContext(ctx) return msg }, expectedDelay: 3 * time.Second, }, { name: "delay with until", publisher: pub, messageConstructor: func(id string) *message.Message { msg := message.NewMessage(id, nil) delay.Message(msg, delay.Until(time.Now().UTC().Add(4*time.Second))) return msg }, expectedDelay: 4 * time.Second, }, { name: "both metadata and context set", publisher: defaultDelayPub, messageConstructor: func(id string) *message.Message { msg := message.NewMessage(id, nil) delay.Message(msg, delay.For(5*time.Second)) ctx := delay.WithContext(context.Background(), delay.For(6*time.Second)) msg.SetContext(ctx) return msg }, expectedDelay: 5 * time.Second, }, } for i, testCase := range testCases { t.Run(testCase.name, func(t *testing.T) { id := fmt.Sprint(i) msg := testCase.messageConstructor(id) err = testCase.publisher.Publish("test", msg) if testCase.expectedError { require.Error(t, err) return } require.NoError(t, err) assertMessage(t, messages, id, testCase.expectedDelay) }) } } func assertMessage(t *testing.T, messages <-chan *message.Message, expectedID string, expectedDelay time.Duration) { t.Helper() select { case msg := <-messages: assert.Equal(t, expectedID, msg.UUID) if expectedDelay == 0 { assert.Empty(t, msg.Metadata.Get(delay.DelayedUntilKey)) assert.Empty(t, msg.Metadata.Get(delay.DelayedForKey)) } else { delayedFor, err := time.ParseDuration(msg.Metadata.Get(delay.DelayedForKey)) require.NoError(t, err) assert.Equal(t, expectedDelay, delayedFor.Round(time.Second)) delayedUntil, err := time.Parse(time.RFC3339, msg.Metadata.Get(delay.DelayedUntilKey)) require.NoError(t, err) assert.WithinDuration(t, time.Now().UTC().Add(expectedDelay), delayedUntil, 1*time.Second) } msg.Ack() case <-time.After(100 * time.Millisecond): require.Fail(t, "timeout") } }