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

     1  package batch
     2  
     3  import (
     4  	"strconv"
     5  	"testing"
     6  
     7  	"github.com/stretchr/testify/require"
     8  
     9  	promchunk "github.com/cortexproject/cortex/pkg/chunk/encoding"
    10  )
    11  
    12  func TestStream(t *testing.T) {
    13  	for i, tc := range []struct {
    14  		input1, input2 []promchunk.Batch
    15  		output         batchStream
    16  	}{
    17  		{
    18  			input1: []promchunk.Batch{mkBatch(0)},
    19  			output: []promchunk.Batch{mkBatch(0)},
    20  		},
    21  
    22  		{
    23  			input1: []promchunk.Batch{mkBatch(0)},
    24  			input2: []promchunk.Batch{mkBatch(0)},
    25  			output: []promchunk.Batch{mkBatch(0)},
    26  		},
    27  
    28  		{
    29  			input1: []promchunk.Batch{mkBatch(0)},
    30  			input2: []promchunk.Batch{mkBatch(promchunk.BatchSize)},
    31  			output: []promchunk.Batch{mkBatch(0), mkBatch(promchunk.BatchSize)},
    32  		},
    33  
    34  		{
    35  			input1: []promchunk.Batch{mkBatch(0), mkBatch(promchunk.BatchSize)},
    36  			input2: []promchunk.Batch{mkBatch(promchunk.BatchSize / 2), mkBatch(2 * promchunk.BatchSize)},
    37  			output: []promchunk.Batch{mkBatch(0), mkBatch(promchunk.BatchSize), mkBatch(2 * promchunk.BatchSize)},
    38  		},
    39  
    40  		{
    41  			input1: []promchunk.Batch{mkBatch(promchunk.BatchSize / 2), mkBatch(3 * promchunk.BatchSize / 2), mkBatch(5 * promchunk.BatchSize / 2)},
    42  			input2: []promchunk.Batch{mkBatch(0), mkBatch(promchunk.BatchSize), mkBatch(3 * promchunk.BatchSize)},
    43  			output: []promchunk.Batch{mkBatch(0), mkBatch(promchunk.BatchSize), mkBatch(2 * promchunk.BatchSize), mkBatch(3 * promchunk.BatchSize)},
    44  		},
    45  	} {
    46  		t.Run(strconv.Itoa(i), func(t *testing.T) {
    47  			result := make(batchStream, len(tc.input1)+len(tc.input2))
    48  			result = mergeStreams(tc.input1, tc.input2, result, promchunk.BatchSize)
    49  			require.Equal(t, batchStream(tc.output), result)
    50  		})
    51  	}
    52  }
    53  
    54  func mkBatch(from int64) promchunk.Batch {
    55  	var result promchunk.Batch
    56  	for i := int64(0); i < promchunk.BatchSize; i++ {
    57  		result.Timestamps[i] = from + i
    58  		result.Values[i] = float64(from + i)
    59  	}
    60  	result.Length = promchunk.BatchSize
    61  	return result
    62  }