github.com/grafana/pyroscope@v1.18.0/pkg/distributor/rate_strategy_test.go (about)

     1  // SPDX-License-Identifier: AGPL-3.0-only
     2  // Provenance-includes-location: https://github.com/cortexproject/cortex/blob/master/pkg/distributor/ingestion_rate_strategy_test.go
     3  // Provenance-includes-license: Apache-2.0
     4  // Provenance-includes-copyright: The Cortex Authors.
     5  
     6  package distributor
     7  
     8  import (
     9  	"testing"
    10  
    11  	"github.com/stretchr/testify/assert"
    12  	"github.com/stretchr/testify/mock"
    13  	"github.com/stretchr/testify/require"
    14  	"golang.org/x/time/rate"
    15  
    16  	"github.com/grafana/pyroscope/pkg/validation"
    17  )
    18  
    19  func TestIngestionRateStrategy(t *testing.T) {
    20  	t.Run("rate limiter should share the limit across the number of distributors", func(t *testing.T) {
    21  		// Init limits overrides
    22  		overrides, err := validation.NewOverrides(validation.Limits{
    23  			IngestionRateMB:      float64(1000),
    24  			IngestionBurstSizeMB: 10000,
    25  		}, nil)
    26  		require.NoError(t, err)
    27  
    28  		mockRing := newReadLifecyclerMock()
    29  		mockRing.On("HealthyInstancesCount").Return(2)
    30  
    31  		strategy := newGlobalRateStrategy(newIngestionRateStrategy(overrides), mockRing)
    32  		assert.Equal(t, strategy.Limit("test"), float64(1000*1024*1024/2))
    33  		assert.Equal(t, strategy.Burst("test"), 10000*1024*1024)
    34  	})
    35  
    36  	t.Run("infinite rate limiter should return unlimited settings", func(t *testing.T) {
    37  		strategy := newInfiniteRateStrategy()
    38  
    39  		assert.Equal(t, strategy.Limit("test"), float64(rate.Inf))
    40  		assert.Equal(t, strategy.Burst("test"), 0)
    41  	})
    42  }
    43  
    44  type readLifecyclerMock struct {
    45  	mock.Mock
    46  }
    47  
    48  func newReadLifecyclerMock() *readLifecyclerMock {
    49  	return &readLifecyclerMock{}
    50  }
    51  
    52  func (m *readLifecyclerMock) HealthyInstancesCount() int {
    53  	args := m.Called()
    54  	return args.Int(0)
    55  }