add unit test

Signed-off-by: Ben Ye <benye@amazon.com>
This commit is contained in:
Ben Ye 2024-06-16 12:09:42 -07:00
parent e7db2e30a4
commit 0e6fca8e76
2 changed files with 29 additions and 2 deletions

View file

@ -1559,8 +1559,10 @@ func (r *Reader) LabelNamesFor(ctx context.Context, postings Postings) ([]string
id := postings.At()
i++
if i%checkContextEveryNIterations == 0 && ctx.Err() != nil {
return nil, ctx.Err()
if i%checkContextEveryNIterations == 0 {
if ctxErr := ctx.Err(); ctxErr != nil {
return nil, ctxErr
}
}
offset := id

View file

@ -634,6 +634,31 @@ func TestReader_PostingsForLabelMatchingHonorsContextCancel(t *testing.T) {
require.Equal(t, failAfter, ctx.Count())
}
func TestReader_LabelNamesForHonorsContextCancel(t *testing.T) {
const seriesCount = 1000
var input indexWriterSeriesSlice
for i := 1; i <= seriesCount; i++ {
input = append(input, &indexWriterSeries{
labels: labels.FromStrings(labels.MetricName, fmt.Sprintf("%4d", i)),
chunks: []chunks.Meta{
{Ref: 1, MinTime: 0, MaxTime: 10},
},
})
}
ir, _, _ := createFileReader(context.Background(), t, input)
name, value := AllPostingsKey()
p, err := ir.Postings(context.Background(), name, value)
require.NoError(t, err)
// We check context cancellation every 128 iterations so 3 will fail after
// iterating 3 * 128 series.
failAfter := uint64(3)
ctx := &testutil.MockContextErrAfter{FailAfter: failAfter}
_, err = ir.LabelNamesFor(ctx, p)
require.Error(t, err)
require.Equal(t, failAfter, ctx.Count())
}
// createFileReader creates a temporary index file. It writes the provided input to this file.
// It returns a Reader for this file, the file's name, and the symbol map.
func createFileReader(ctx context.Context, tb testing.TB, input indexWriterSeriesSlice) (*Reader, string, map[string]struct{}) {