github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/pkg/querier/batch/chunk_test.go (about)

     1  package batch
     2  
     3  import (
     4  	"strconv"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/prometheus/common/model"
     9  	"github.com/prometheus/prometheus/pkg/labels"
    10  	"github.com/prometheus/prometheus/tsdb/chunkenc"
    11  	"github.com/stretchr/testify/require"
    12  
    13  	"github.com/cortexproject/cortex/pkg/chunk"
    14  	promchunk "github.com/cortexproject/cortex/pkg/chunk/encoding"
    15  )
    16  
    17  const (
    18  	fp     = 1
    19  	userID = "0"
    20  	step   = 1 * time.Second
    21  )
    22  
    23  func TestChunkIter(t *testing.T) {
    24  	forEncodings(t, func(t *testing.T, enc promchunk.Encoding) {
    25  		chunk := mkGenericChunk(t, 0, 100, enc)
    26  		iter := &chunkIterator{}
    27  
    28  		iter.reset(chunk)
    29  		testIter(t, 100, newIteratorAdapter(iter))
    30  
    31  		iter.reset(chunk)
    32  		testSeek(t, 100, newIteratorAdapter(iter))
    33  	})
    34  }
    35  
    36  func forEncodings(t *testing.T, f func(t *testing.T, enc promchunk.Encoding)) {
    37  	for _, enc := range []promchunk.Encoding{
    38  		promchunk.DoubleDelta, promchunk.Varbit, promchunk.Bigchunk, promchunk.PrometheusXorChunk,
    39  	} {
    40  		t.Run(enc.String(), func(t *testing.T) {
    41  			f(t, enc)
    42  		})
    43  	}
    44  }
    45  
    46  func mkChunk(t require.TestingT, from model.Time, points int, enc promchunk.Encoding) chunk.Chunk {
    47  	metric := labels.Labels{
    48  		{Name: model.MetricNameLabel, Value: "foo"},
    49  	}
    50  	pc, err := promchunk.NewForEncoding(enc)
    51  	require.NoError(t, err)
    52  	ts := from
    53  	for i := 0; i < points; i++ {
    54  		npc, err := pc.Add(model.SamplePair{
    55  			Timestamp: ts,
    56  			Value:     model.SampleValue(float64(ts)),
    57  		})
    58  		require.NoError(t, err)
    59  		require.Nil(t, npc)
    60  		ts = ts.Add(step)
    61  	}
    62  	ts = ts.Add(-step) // undo the add that we did just before exiting the loop
    63  	return chunk.NewChunk(userID, fp, metric, pc, from, ts)
    64  }
    65  
    66  func mkGenericChunk(t require.TestingT, from model.Time, points int, enc promchunk.Encoding) GenericChunk {
    67  	ck := mkChunk(t, from, points, enc)
    68  	return NewGenericChunk(int64(ck.From), int64(ck.Through), ck.Data.NewIterator)
    69  }
    70  
    71  func testIter(t require.TestingT, points int, iter chunkenc.Iterator) {
    72  	ets := model.TimeFromUnix(0)
    73  	for i := 0; i < points; i++ {
    74  		require.True(t, iter.Next(), strconv.Itoa(i))
    75  		ts, v := iter.At()
    76  		require.EqualValues(t, int64(ets), ts, strconv.Itoa(i))
    77  		require.EqualValues(t, float64(ets), v, strconv.Itoa(i))
    78  		ets = ets.Add(step)
    79  	}
    80  	require.False(t, iter.Next())
    81  }
    82  
    83  func testSeek(t require.TestingT, points int, iter chunkenc.Iterator) {
    84  	for i := 0; i < points; i += points / 10 {
    85  		ets := int64(i * int(step/time.Millisecond))
    86  
    87  		require.True(t, iter.Seek(ets))
    88  		ts, v := iter.At()
    89  		require.EqualValues(t, ets, ts)
    90  		require.EqualValues(t, v, float64(ets))
    91  		require.NoError(t, iter.Err())
    92  
    93  		for j := i + 1; j < i+points/10; j++ {
    94  			ets := int64(j * int(step/time.Millisecond))
    95  			require.True(t, iter.Next())
    96  			ts, v := iter.At()
    97  			require.EqualValues(t, ets, ts)
    98  			require.EqualValues(t, float64(ets), v)
    99  			require.NoError(t, iter.Err())
   100  		}
   101  	}
   102  }
   103  
   104  func TestSeek(t *testing.T) {
   105  	var it mockIterator
   106  	c := chunkIterator{
   107  		chunk: GenericChunk{
   108  			MaxTime: promchunk.BatchSize,
   109  		},
   110  		it: &it,
   111  	}
   112  
   113  	for i := 0; i < promchunk.BatchSize-1; i++ {
   114  		require.True(t, c.Seek(int64(i), 1))
   115  	}
   116  	require.Equal(t, 1, it.seeks)
   117  
   118  	require.True(t, c.Seek(int64(promchunk.BatchSize), 1))
   119  	require.Equal(t, 2, it.seeks)
   120  }
   121  
   122  type mockIterator struct {
   123  	seeks int
   124  }
   125  
   126  func (i *mockIterator) Scan() bool {
   127  	return true
   128  }
   129  
   130  func (i *mockIterator) FindAtOrAfter(model.Time) bool {
   131  	i.seeks++
   132  	return true
   133  }
   134  
   135  func (i *mockIterator) Value() model.SamplePair {
   136  	return model.SamplePair{}
   137  }
   138  
   139  func (i *mockIterator) Batch(size int) promchunk.Batch {
   140  	batch := promchunk.Batch{
   141  		Length: promchunk.BatchSize,
   142  	}
   143  	for i := 0; i < promchunk.BatchSize; i++ {
   144  		batch.Timestamps[i] = int64(i)
   145  	}
   146  	return batch
   147  }
   148  
   149  func (i *mockIterator) Err() error {
   150  	return nil
   151  }