github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/pkg/util/limiter/query_limiter_test.go (about)

     1  package limiter
     2  
     3  import (
     4  	"fmt"
     5  	"testing"
     6  
     7  	"github.com/prometheus/prometheus/pkg/labels"
     8  	"github.com/stretchr/testify/assert"
     9  	"github.com/stretchr/testify/require"
    10  
    11  	"github.com/cortexproject/cortex/pkg/cortexpb"
    12  )
    13  
    14  func TestQueryLimiter_AddSeries_ShouldReturnNoErrorOnLimitNotExceeded(t *testing.T) {
    15  	const (
    16  		metricName = "test_metric"
    17  	)
    18  
    19  	var (
    20  		series1 = labels.FromMap(map[string]string{
    21  			labels.MetricName: metricName + "_1",
    22  			"series1":         "1",
    23  		})
    24  		series2 = labels.FromMap(map[string]string{
    25  			labels.MetricName: metricName + "_2",
    26  			"series2":         "1",
    27  		})
    28  		limiter = NewQueryLimiter(100, 0, 0)
    29  	)
    30  	err := limiter.AddSeries(cortexpb.FromLabelsToLabelAdapters(series1))
    31  	assert.NoError(t, err)
    32  	err = limiter.AddSeries(cortexpb.FromLabelsToLabelAdapters(series2))
    33  	assert.NoError(t, err)
    34  	assert.Equal(t, 2, limiter.uniqueSeriesCount())
    35  
    36  	// Re-add previous series to make sure it's not double counted
    37  	err = limiter.AddSeries(cortexpb.FromLabelsToLabelAdapters(series1))
    38  	assert.NoError(t, err)
    39  	assert.Equal(t, 2, limiter.uniqueSeriesCount())
    40  }
    41  
    42  func TestQueryLimiter_AddSeriers_ShouldReturnErrorOnLimitExceeded(t *testing.T) {
    43  	const (
    44  		metricName = "test_metric"
    45  	)
    46  
    47  	var (
    48  		series1 = labels.FromMap(map[string]string{
    49  			labels.MetricName: metricName + "_1",
    50  			"series1":         "1",
    51  		})
    52  		series2 = labels.FromMap(map[string]string{
    53  			labels.MetricName: metricName + "_2",
    54  			"series2":         "1",
    55  		})
    56  		limiter = NewQueryLimiter(1, 0, 0)
    57  	)
    58  	err := limiter.AddSeries(cortexpb.FromLabelsToLabelAdapters(series1))
    59  	require.NoError(t, err)
    60  	err = limiter.AddSeries(cortexpb.FromLabelsToLabelAdapters(series2))
    61  	require.Error(t, err)
    62  }
    63  
    64  func TestQueryLimiter_AddChunkBytes(t *testing.T) {
    65  	var limiter = NewQueryLimiter(0, 100, 0)
    66  
    67  	err := limiter.AddChunkBytes(100)
    68  	require.NoError(t, err)
    69  	err = limiter.AddChunkBytes(1)
    70  	require.Error(t, err)
    71  }
    72  
    73  func BenchmarkQueryLimiter_AddSeries(b *testing.B) {
    74  	const (
    75  		metricName = "test_metric"
    76  	)
    77  	var series []labels.Labels
    78  	for i := 0; i < b.N; i++ {
    79  		series = append(series,
    80  			labels.FromMap(map[string]string{
    81  				labels.MetricName: metricName + "_1",
    82  				"series1":         fmt.Sprint(i),
    83  			}))
    84  	}
    85  	b.ResetTimer()
    86  
    87  	limiter := NewQueryLimiter(b.N+1, 0, 0)
    88  	for _, s := range series {
    89  		err := limiter.AddSeries(cortexpb.FromLabelsToLabelAdapters(s))
    90  		assert.NoError(b, err)
    91  	}
    92  
    93  }