github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/querier/queryrange/queryrangebase/test_utils_test.go (about)

     1  package queryrangebase
     2  
     3  import (
     4  	"math"
     5  	"sort"
     6  	"testing"
     7  
     8  	"github.com/prometheus/prometheus/model/labels"
     9  	"github.com/stretchr/testify/require"
    10  
    11  	"github.com/grafana/loki/pkg/querier/astmapper"
    12  )
    13  
    14  func TestGenLabelsCorrectness(t *testing.T) {
    15  	ls := genLabels([]string{"a", "b"}, 2)
    16  	for _, set := range ls {
    17  		sort.Sort(set)
    18  	}
    19  	expected := []labels.Labels{
    20  		{
    21  			labels.Label{
    22  				Name:  "a",
    23  				Value: "0",
    24  			},
    25  			labels.Label{
    26  				Name:  "b",
    27  				Value: "0",
    28  			},
    29  		},
    30  		{
    31  			labels.Label{
    32  				Name:  "a",
    33  				Value: "0",
    34  			},
    35  			labels.Label{
    36  				Name:  "b",
    37  				Value: "1",
    38  			},
    39  		},
    40  		{
    41  			labels.Label{
    42  				Name:  "a",
    43  				Value: "1",
    44  			},
    45  			labels.Label{
    46  				Name:  "b",
    47  				Value: "0",
    48  			},
    49  		},
    50  		{
    51  			labels.Label{
    52  				Name:  "a",
    53  				Value: "1",
    54  			},
    55  			labels.Label{
    56  				Name:  "b",
    57  				Value: "1",
    58  			},
    59  		},
    60  	}
    61  	require.Equal(t, expected, ls)
    62  }
    63  
    64  func TestGenLabelsSize(t *testing.T) {
    65  	for _, tc := range []struct {
    66  		set     []string
    67  		buckets int
    68  	}{
    69  		{
    70  			set:     []string{"a", "b"},
    71  			buckets: 5,
    72  		},
    73  		{
    74  			set:     []string{"a", "b", "c"},
    75  			buckets: 10,
    76  		},
    77  	} {
    78  		sets := genLabels(tc.set, tc.buckets)
    79  		require.Equal(
    80  			t,
    81  			math.Pow(float64(tc.buckets), float64(len(tc.set))),
    82  			float64(len(sets)),
    83  		)
    84  	}
    85  }
    86  
    87  func TestNewMockShardedqueryable(t *testing.T) {
    88  	for _, tc := range []struct {
    89  		shards, nSamples, labelBuckets int
    90  		labelSet                       []string
    91  	}{
    92  		{
    93  			nSamples:     100,
    94  			shards:       1,
    95  			labelBuckets: 3,
    96  			labelSet:     []string{"a", "b", "c"},
    97  		},
    98  		{
    99  			nSamples:     0,
   100  			shards:       2,
   101  			labelBuckets: 3,
   102  			labelSet:     []string{"a", "b", "c"},
   103  		},
   104  	} {
   105  		q := NewMockShardedQueryable(tc.nSamples, tc.labelSet, tc.labelBuckets, 0)
   106  		expectedSeries := int(math.Pow(float64(tc.labelBuckets), float64(len(tc.labelSet))))
   107  
   108  		seriesCt := 0
   109  		for i := 0; i < tc.shards; i++ {
   110  
   111  			set := q.Select(false, nil, &labels.Matcher{
   112  				Type: labels.MatchEqual,
   113  				Name: astmapper.ShardLabel,
   114  				Value: astmapper.ShardAnnotation{
   115  					Shard: i,
   116  					Of:    tc.shards,
   117  				}.String(),
   118  			})
   119  
   120  			require.Nil(t, set.Err())
   121  
   122  			for set.Next() {
   123  				seriesCt++
   124  				iter := set.At().Iterator()
   125  				samples := 0
   126  				for iter.Next() {
   127  					samples++
   128  				}
   129  				require.Equal(t, tc.nSamples, samples)
   130  			}
   131  
   132  		}
   133  		require.Equal(t, expectedSeries, seriesCt)
   134  	}
   135  }