feat: use defaul rate limiter when nil

waku-org · Dec 5, 2024 · 76309b2 · 76309b2
1 parent b6cf713
commit 76309b2
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 7 deletions.
diff --git a/waku/v2/api/publish/message_sender.go b/waku/v2/api/publish/message_sender.go
@@ -79,14 +79,19 @@ func (r *Request) WithPublishMethod(publishMethod PublishMethod) *Request {
 	return r
 }
 
-func NewMessageSender(publishMethod PublishMethod, publisher Publisher, logger *zap.Logger) (*MessageSender, error) {
+func NewMessageSender(publishMethod PublishMethod, publisher Publisher, rateLimiter PublishRateLimiter, logger *zap.Logger) (*MessageSender, error) {
 	if publishMethod == UnknownMethod {
 		return nil, errors.New("publish method is required")
 	}
+
+	if rateLimiter == nil {
+		rateLimiter = NewDefaultRateLimiter(DefaultPublishingLimiterRate, DefaultPublishingLimitBurst)
+	}
+
 	return &MessageSender{
 		publishMethod: publishMethod,
 		publisher:     publisher,
-		rateLimiter:   NewRlnRateLimiter(RlnLimiterCapacity, RlnLimiterRefillInterval),
+		rateLimiter:   rateLimiter,
 		logger:        logger,
 	}, nil
 }

diff --git a/waku/v2/api/publish/message_sender_test.go b/waku/v2/api/publish/message_sender_test.go
@@ -41,7 +41,7 @@ func (m *MockMessageSentCheck) Start() {
 }
 
 func TestNewSenderWithUnknownMethod(t *testing.T) {
-	sender, err := NewMessageSender(UnknownMethod, nil, nil)
+	sender, err := NewMessageSender(UnknownMethod, nil, nil, nil)
 	require.NotNil(t, err)
 	require.Nil(t, sender)
 }
@@ -55,7 +55,7 @@ func TestNewSenderWithRelay(t *testing.T) {
 	_, err = relayNode.Subscribe(context.Background(), protocol.NewContentFilter("test-pubsub-topic"))
 	require.Nil(t, err)
 	publisher := NewDefaultPublisher(nil, relayNode)
-	sender, err := NewMessageSender(Relay, publisher, utils.Logger())
+	sender, err := NewMessageSender(Relay, publisher, nil, utils.Logger())
 	require.Nil(t, err)
 	require.NotNil(t, sender)
 	require.Nil(t, sender.messageSentCheck)
@@ -81,7 +81,7 @@ func TestNewSenderWithRelayAndMessageSentCheck(t *testing.T) {
 	_, err = relayNode.Subscribe(context.Background(), protocol.NewContentFilter("test-pubsub-topic"))
 	require.Nil(t, err)
 	publisher := NewDefaultPublisher(nil, relayNode)
-	sender, err := NewMessageSender(Relay, publisher, utils.Logger())
+	sender, err := NewMessageSender(Relay, publisher, nil, utils.Logger())
 
 	check := &MockMessageSentCheck{Messages: make(map[string]map[common.Hash]uint32)}
 	sender.WithMessageSentCheck(check)
@@ -111,7 +111,7 @@ func TestNewSenderWithRelayAndMessageSentCheck(t *testing.T) {
 }
 
 func TestNewSenderWithLightPush(t *testing.T) {
-	sender, err := NewMessageSender(LightPush, nil, nil)
+	sender, err := NewMessageSender(LightPush, nil, nil, nil)
 	require.Nil(t, err)
 	require.NotNil(t, sender)
 	require.Equal(t, LightPush, sender.publishMethod)
@@ -140,7 +140,7 @@ func TestMessageSentEmitter(t *testing.T) {
 	_, err = relayNode.Subscribe(context.Background(), protocol.NewContentFilter("test-pubsub-topic"))
 	require.Nil(t, err)
 	publisher := NewDefaultPublisher(nil, relayNode)
-	sender, err := NewMessageSender(Relay, publisher, utils.Logger())
+	sender, err := NewMessageSender(Relay, publisher, nil, utils.Logger())
 	require.Nil(t, err)
 
 	check := &MockMessageSentCheck{Messages: make(map[string]map[common.Hash]uint32)}