github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/pkg/alertmanager/rate_limited_notifier_test.go (about)

     1  package alertmanager
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/prometheus/alertmanager/types"
     9  	"github.com/prometheus/client_golang/prometheus"
    10  	"github.com/prometheus/client_golang/prometheus/testutil"
    11  	"github.com/stretchr/testify/assert"
    12  	"golang.org/x/time/rate"
    13  )
    14  
    15  func TestRateLimitedNotifier(t *testing.T) {
    16  	mock := &mockNotifier{}
    17  	counter := prometheus.NewCounter(prometheus.CounterOpts{})
    18  
    19  	// Initial limits.
    20  	limiter := &limiter{limit: 5, burst: 5}
    21  	rateLimitedNotifier := newRateLimitedNotifier(mock, limiter, 10*time.Second, counter)
    22  
    23  	runNotifications(t, rateLimitedNotifier, counter, 10, 5, 5, 5)
    24  
    25  	// Disable limits.
    26  	limiter.limit = rate.Inf
    27  	limiter.burst = 0
    28  
    29  	runNotifications(t, rateLimitedNotifier, counter, 10, 10, 0, 5)
    30  
    31  	limiter.limit = 5
    32  	limiter.burst = 5
    33  
    34  	// We don't get any successful notifications until some time passes.
    35  	runNotifications(t, rateLimitedNotifier, counter, 10, 0, 10, 15)
    36  
    37  	time.Sleep(1 * time.Second) // Wait to refill rate-limiter's "bucket".
    38  	runNotifications(t, rateLimitedNotifier, counter, 10, 5, 5, 20)
    39  }
    40  
    41  func runNotifications(t *testing.T, rateLimitedNotifier *rateLimitedNotifier, counter prometheus.Counter, count, expectedSuccess, expectedRateLimited, expectedCounter int) {
    42  	rateLimitedNotifier.recheckAt.Store(0) // Force recheck of limits.
    43  
    44  	success := 0
    45  	rateLimited := 0
    46  
    47  	for i := 0; i < count; i++ {
    48  		retry, err := rateLimitedNotifier.Notify(context.Background(), &types.Alert{})
    49  
    50  		if err == nil {
    51  			success++
    52  		} else if err == errRateLimited {
    53  			rateLimited++
    54  			assert.False(t, retry)
    55  		} else {
    56  			assert.NotNil(t, err)
    57  		}
    58  	}
    59  
    60  	assert.Equal(t, expectedSuccess, success)
    61  	assert.Equal(t, expectedRateLimited, rateLimited)
    62  	assert.Equal(t, expectedCounter, int(testutil.ToFloat64(counter)))
    63  }
    64  
    65  type mockNotifier struct{}
    66  
    67  func (m *mockNotifier) Notify(ctx context.Context, alert ...*types.Alert) (bool, error) {
    68  	return false, nil
    69  }
    70  
    71  type limiter struct {
    72  	limit rate.Limit
    73  	burst int
    74  }
    75  
    76  func (l *limiter) RateLimit() rate.Limit {
    77  	return l.limit
    78  }
    79  
    80  func (l *limiter) Burst() int {
    81  	return l.burst
    82  }