github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/storage/stores/shipper/indexgateway/gateway_test.go (about)

     1  package indexgateway
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"testing"
     7  
     8  	"github.com/stretchr/testify/require"
     9  	"google.golang.org/grpc"
    10  
    11  	"github.com/grafana/loki/pkg/logproto"
    12  	"github.com/grafana/loki/pkg/storage/stores/series/index"
    13  	"github.com/grafana/loki/pkg/storage/stores/shipper/util"
    14  	util_math "github.com/grafana/loki/pkg/util/math"
    15  )
    16  
    17  const (
    18  	// query prefixes
    19  	tableNamePrefix        = "table-name"
    20  	hashValuePrefix        = "hash-value"
    21  	rangeValuePrefixPrefix = "range-value-prefix"
    22  	rangeValueStartPrefix  = "range-value-start"
    23  	valueEqualPrefix       = "value-equal"
    24  
    25  	// response prefixes
    26  	rangeValuePrefix = "range-value"
    27  	valuePrefix      = "value"
    28  )
    29  
    30  type mockBatch struct {
    31  	size int
    32  }
    33  
    34  func (r *mockBatch) Iterator() index.ReadBatchIterator {
    35  	return &mockBatchIter{
    36  		curr: -1,
    37  		size: r.size,
    38  	}
    39  }
    40  
    41  type mockBatchIter struct {
    42  	curr, size int
    43  }
    44  
    45  func (b *mockBatchIter) Next() bool {
    46  	b.curr++
    47  	return b.curr < b.size
    48  }
    49  
    50  func (b *mockBatchIter) RangeValue() []byte {
    51  	return []byte(fmt.Sprintf("%s%d", rangeValuePrefix, b.curr))
    52  }
    53  
    54  func (b *mockBatchIter) Value() []byte {
    55  	return []byte(fmt.Sprintf("%s%d", valuePrefix, b.curr))
    56  }
    57  
    58  type mockQueryIndexServer struct {
    59  	grpc.ServerStream
    60  	callback func(resp *logproto.QueryIndexResponse)
    61  }
    62  
    63  func (m *mockQueryIndexServer) Send(resp *logproto.QueryIndexResponse) error {
    64  	m.callback(resp)
    65  	return nil
    66  }
    67  
    68  func (m *mockQueryIndexServer) Context() context.Context {
    69  	return context.Background()
    70  }
    71  
    72  type mockIndexClient struct {
    73  	index.Client
    74  	response *mockBatch
    75  }
    76  
    77  func (m mockIndexClient) QueryPages(ctx context.Context, queries []index.Query, callback index.QueryPagesCallback) error {
    78  	for _, query := range queries {
    79  		callback(query, m.response)
    80  	}
    81  
    82  	return nil
    83  }
    84  
    85  func TestGateway_QueryIndex(t *testing.T) {
    86  	var expectedQueryKey string
    87  	type batchRange struct {
    88  		start, end int
    89  	}
    90  	var expectedRanges []batchRange
    91  
    92  	// the response should have index entries between start and end from batchRange at index 0 of expectedRanges
    93  	var server logproto.IndexGateway_QueryIndexServer = &mockQueryIndexServer{
    94  		callback: func(resp *logproto.QueryIndexResponse) {
    95  			require.Equal(t, expectedQueryKey, resp.QueryKey)
    96  
    97  			require.True(t, len(expectedRanges) > 0)
    98  			require.Len(t, resp.Rows, expectedRanges[0].end-expectedRanges[0].start)
    99  			i := expectedRanges[0].start
   100  			for _, row := range resp.Rows {
   101  				require.Equal(t, fmt.Sprintf("%s%d", rangeValuePrefix, i), string(row.RangeValue))
   102  				require.Equal(t, fmt.Sprintf("%s%d", valuePrefix, i), string(row.Value))
   103  				i++
   104  			}
   105  
   106  			// remove first element for checking the response from next callback.
   107  			expectedRanges = expectedRanges[1:]
   108  		},
   109  	}
   110  
   111  	gateway := Gateway{}
   112  	responseSizes := []int{0, 99, maxIndexEntriesPerResponse, 2 * maxIndexEntriesPerResponse, 5*maxIndexEntriesPerResponse - 1}
   113  	for i, responseSize := range responseSizes {
   114  		query := index.Query{
   115  			TableName:        fmt.Sprintf("%s%d", tableNamePrefix, i),
   116  			HashValue:        fmt.Sprintf("%s%d", hashValuePrefix, i),
   117  			RangeValuePrefix: []byte(fmt.Sprintf("%s%d", rangeValuePrefixPrefix, i)),
   118  			RangeValueStart:  []byte(fmt.Sprintf("%s%d", rangeValueStartPrefix, i)),
   119  			ValueEqual:       []byte(fmt.Sprintf("%s%d", valueEqualPrefix, i)),
   120  		}
   121  
   122  		// build expectedRanges based on maxIndexEntriesPerResponse
   123  		for j := 0; j < responseSize; j += maxIndexEntriesPerResponse {
   124  			expectedRanges = append(expectedRanges, batchRange{
   125  				start: j,
   126  				end:   util_math.Min(j+maxIndexEntriesPerResponse, responseSize),
   127  			})
   128  		}
   129  		expectedQueryKey = util.QueryKey(query)
   130  		gateway.indexClient = mockIndexClient{response: &mockBatch{size: responseSize}}
   131  
   132  		err := gateway.QueryIndex(&logproto.QueryIndexRequest{Queries: []*logproto.IndexQuery{{
   133  			TableName:        query.TableName,
   134  			HashValue:        query.HashValue,
   135  			RangeValuePrefix: query.RangeValuePrefix,
   136  			RangeValueStart:  query.RangeValueStart,
   137  			ValueEqual:       query.ValueEqual,
   138  		}}}, server)
   139  		require.NoError(t, err)
   140  
   141  		// verify that we actually got responses back by checking if expectedRanges got cleared.
   142  		require.Len(t, expectedRanges, 0)
   143  	}
   144  }