add more tests in ut

Signed-off-by: Ben Ye <[email protected]>

Author: yeya24

integration/querier_remote_read_test.go

@@ -6,6 +6,7 @@ package integration
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"io"
 	"net/http"
 	"strconv"
@@ -26,21 +27,14 @@ import (
 )
 
 func TestQuerierRemoteRead(t *testing.T) {
-	tests := map[string]struct {
-		thanosEngine bool
-	}{
-		"old engine": {thanosEngine: false},
-		"new engine": {thanosEngine: true},
-	}
-
-	for testName, tc := range tests {
-		t.Run(testName, func(t *testing.T) {
+	for _, thanosEngine := range []bool{false, true} {
+		t.Run(fmt.Sprintf("thanosEngine=%t", thanosEngine), func(t *testing.T) {
 			s, err := e2e.NewScenario(networkName)
 			require.NoError(t, err)
 			defer s.Close()
 
 			flags := mergeFlags(BlocksStorageFlags(), map[string]string{
-				"-querier.thanos-engine": strconv.FormatBool(tc.thanosEngine),
+				"-querier.thanos-engine": strconv.FormatBool(thanosEngine),
 			})
 
 			// Start dependencies.

integration/querier_test.go

@@ -39,14 +39,8 @@ func TestRulerAPI(t *testing.T) {
 	)
 	ruleGroup := createTestRuleGroup(t)
 
-	tests := map[string]struct {
-		thanosEngine bool
-	}{
-		"old engine": {thanosEngine: false},
-		"new engine": {thanosEngine: true},
-	}
-	for testName, tc := range tests {
-		t.Run(testName, func(t *testing.T) {
+	for _, thanosEngine := range []bool{false, true} {
+		t.Run(fmt.Sprintf("thanosEngine=%t", thanosEngine), func(t *testing.T) {
 			s, err := e2e.NewScenario(networkName)
 			require.NoError(t, err)
 			defer s.Close()
@@ -57,7 +51,7 @@ func TestRulerAPI(t *testing.T) {
 			require.NoError(t, s.StartAndWaitReady(consul, minio))
 
 			flags := mergeFlags(BlocksStorageFlags(), RulerFlags(), map[string]string{
-				"-querier.thanos-engine": strconv.FormatBool(tc.thanosEngine),
+				"-querier.thanos-engine": strconv.FormatBool(thanosEngine),
 			})
 
 			// Start Cortex components.

integration/ruler_test.go

@@ -22,9 +22,12 @@ import (
 	"github.com/prometheus/prometheus/promql"
 	"github.com/prometheus/prometheus/storage"
 	"github.com/prometheus/prometheus/tsdb/chunkenc"
+	v1 "github.com/prometheus/prometheus/web/api/v1"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
+	"github.com/thanos-community/promql-engine/engine"
+	"github.com/thanos-community/promql-engine/logicalplan"
 	"github.com/thanos-io/thanos/pkg/store/hintspb"
 	"github.com/thanos-io/thanos/pkg/store/labelpb"
 	"github.com/thanos-io/thanos/pkg/store/storepb"
@@ -1319,6 +1322,13 @@ func TestBlocksStoreQuerier_SelectSortedShouldHonorQueryStoreAfter(t *testing.T)
 }
 
 func TestBlocksStoreQuerier_PromQLExecution(t *testing.T) {
+	logger := log.NewNopLogger()
+	opts := promql.EngineOpts{
+		Logger:     logger,
+		Timeout:    10 * time.Second,
+		MaxSamples: 1e6,
+	}
+
 	block1 := ulid.MustNew(1, nil)
 	block2 := ulid.MustNew(2, nil)
 	series1 := []labelpb.ZLabel{{Name: "__name__", Value: "metric_1"}}
@@ -1341,103 +1351,109 @@ func TestBlocksStoreQuerier_PromQLExecution(t *testing.T) {
 		{T: 1589760015000, V: 2},
 		{T: 1589760030000, V: 2},
 	}
+	for _, thanosEngine := range []bool{false, true} {
+		var queryEngine v1.QueryEngine
+		if thanosEngine {
+			queryEngine = engine.New(engine.Opts{
+				EngineOpts:        opts,
+				LogicalOptimizers: logicalplan.AllOptimizers,
+			})
+		} else {
+			queryEngine = promql.NewEngine(opts)
+		}
 
-	// Mock the finder to simulate we need to query two blocks.
-	finder := &blocksFinderMock{
-		Service: services.NewIdleService(nil, nil),
-	}
-	finder.On("GetBlocks", mock.Anything, "user-1", mock.Anything, mock.Anything).Return(bucketindex.Blocks{
-		{ID: block1},
-		{ID: block2},
-	}, map[ulid.ULID]*bucketindex.BlockDeletionMark(nil), error(nil))
-
-	// Mock the store to simulate each block is queried from a different store-gateway.
-	gateway1 := &storeGatewayClientMock{remoteAddr: "1.1.1.1", mockedSeriesResponses: []*storepb.SeriesResponse{
-		{
-			Result: &storepb.SeriesResponse_Series{
-				Series: &storepb.Series{
-					Labels: series1,
-					Chunks: []storepb.AggrChunk{
-						createAggrChunkWithSamples(series1Samples[:3]...), // First half.
+		t.Run(fmt.Sprintf("thanos engine enabled=%t", thanosEngine), func(t *testing.T) {
+			// Mock the finder to simulate we need to query two blocks.
+			finder := &blocksFinderMock{
+				Service: services.NewIdleService(nil, nil),
+			}
+			finder.On("GetBlocks", mock.Anything, "user-1", mock.Anything, mock.Anything).Return(bucketindex.Blocks{
+				{ID: block1},
+				{ID: block2},
+			}, map[ulid.ULID]*bucketindex.BlockDeletionMark(nil), error(nil))
+
+			// Mock the store to simulate each block is queried from a different store-gateway.
+			gateway1 := &storeGatewayClientMock{remoteAddr: "1.1.1.1", mockedSeriesResponses: []*storepb.SeriesResponse{
+				{
+					Result: &storepb.SeriesResponse_Series{
+						Series: &storepb.Series{
+							Labels: series1,
+							Chunks: []storepb.AggrChunk{
+								createAggrChunkWithSamples(series1Samples[:3]...), // First half.
+							},
+						},
 					},
-				},
-			},
-		}, {
-			Result: &storepb.SeriesResponse_Series{
-				Series: &storepb.Series{
-					Labels: series2,
-					Chunks: []storepb.AggrChunk{
-						createAggrChunkWithSamples(series2Samples[:3]...),
+				}, {
+					Result: &storepb.SeriesResponse_Series{
+						Series: &storepb.Series{
+							Labels: series2,
+							Chunks: []storepb.AggrChunk{
+								createAggrChunkWithSamples(series2Samples[:3]...),
+							},
+						},
 					},
 				},
-			},
-		},
-		mockHintsResponse(block1),
-	}}
-
-	gateway2 := &storeGatewayClientMock{remoteAddr: "2.2.2.2", mockedSeriesResponses: []*storepb.SeriesResponse{
-		{
-			Result: &storepb.SeriesResponse_Series{
-				Series: &storepb.Series{
-					Labels: series1,
-					Chunks: []storepb.AggrChunk{
-						createAggrChunkWithSamples(series1Samples[3:]...), // Second half.
+				mockHintsResponse(block1),
+			}}
+
+			gateway2 := &storeGatewayClientMock{remoteAddr: "2.2.2.2", mockedSeriesResponses: []*storepb.SeriesResponse{
+				{
+					Result: &storepb.SeriesResponse_Series{
+						Series: &storepb.Series{
+							Labels: series1,
+							Chunks: []storepb.AggrChunk{
+								createAggrChunkWithSamples(series1Samples[3:]...), // Second half.
+							},
+						},
+					},
+				}, {
+					Result: &storepb.SeriesResponse_Series{
+						Series: &storepb.Series{
+							Labels: series2,
+							Chunks: []storepb.AggrChunk{
+								createAggrChunkWithSamples(series2Samples[3:]...),
+							},
+						},
 					},
 				},
-			},
-		}, {
-			Result: &storepb.SeriesResponse_Series{
-				Series: &storepb.Series{
-					Labels: series2,
-					Chunks: []storepb.AggrChunk{
-						createAggrChunkWithSamples(series2Samples[3:]...),
+				mockHintsResponse(block2),
+			}}
+
+			stores := &blocksStoreSetMock{
+				Service: services.NewIdleService(nil, nil),
+				mockedResponses: []interface{}{
+					map[BlocksStoreClient][]ulid.ULID{
+						gateway1: {block1},
+						gateway2: {block2},
 					},
 				},
-			},
-		},
-		mockHintsResponse(block2),
-	}}
+			}
 
-	stores := &blocksStoreSetMock{
-		Service: services.NewIdleService(nil, nil),
-		mockedResponses: []interface{}{
-			map[BlocksStoreClient][]ulid.ULID{
-				gateway1: {block1},
-				gateway2: {block2},
-			},
-		},
+			// Instance the querier that will be executed to run the query.
+			queryable, err := NewBlocksStoreQueryable(stores, finder, NewBlocksConsistencyChecker(0, 0, logger, nil), &blocksStoreLimitsMock{}, 0, logger, nil)
+			require.NoError(t, err)
+			require.NoError(t, services.StartAndAwaitRunning(context.Background(), queryable))
+			defer services.StopAndAwaitTerminated(context.Background(), queryable) // nolint:errcheck
+
+			// Run a query.
+			q, err := queryEngine.NewRangeQuery(queryable, nil, `{__name__=~"metric.*"}`, time.Unix(1589759955, 0), time.Unix(1589760030, 0), 15*time.Second)
+			require.NoError(t, err)
+
+			ctx := user.InjectOrgID(context.Background(), "user-1")
+			res := q.Exec(ctx)
+			require.NoError(t, err)
+			require.NoError(t, res.Err)
+
+			matrix, err := res.Matrix()
+			require.NoError(t, err)
+			require.Len(t, matrix, 2)
+
+			assert.Equal(t, labelpb.ZLabelsToPromLabels(series1), matrix[0].Metric)
+			assert.Equal(t, labelpb.ZLabelsToPromLabels(series2), matrix[1].Metric)
+			assert.Equal(t, series1Samples, matrix[0].Points)
+			assert.Equal(t, series2Samples, matrix[1].Points)
+		})
 	}
-
-	// Instance the querier that will be executed to run the query.
-	logger := log.NewNopLogger()
-	queryable, err := NewBlocksStoreQueryable(stores, finder, NewBlocksConsistencyChecker(0, 0, logger, nil), &blocksStoreLimitsMock{}, 0, logger, nil)
-	require.NoError(t, err)
-	require.NoError(t, services.StartAndAwaitRunning(context.Background(), queryable))
-	defer services.StopAndAwaitTerminated(context.Background(), queryable) // nolint:errcheck
-
-	engine := promql.NewEngine(promql.EngineOpts{
-		Logger:     logger,
-		Timeout:    10 * time.Second,
-		MaxSamples: 1e6,
-	})
-
-	// Run a query.
-	q, err := engine.NewRangeQuery(queryable, nil, `{__name__=~"metric.*"}`, time.Unix(1589759955, 0), time.Unix(1589760030, 0), 15*time.Second)
-	require.NoError(t, err)
-
-	ctx := user.InjectOrgID(context.Background(), "user-1")
-	res := q.Exec(ctx)
-	require.NoError(t, err)
-	require.NoError(t, res.Err)
-
-	matrix, err := res.Matrix()
-	require.NoError(t, err)
-	require.Len(t, matrix, 2)
-
-	assert.Equal(t, labelpb.ZLabelsToPromLabels(series1), matrix[0].Metric)
-	assert.Equal(t, labelpb.ZLabelsToPromLabels(series2), matrix[1].Metric)
-	assert.Equal(t, series1Samples, matrix[0].Points)
-	assert.Equal(t, series2Samples, matrix[1].Points)
 }
 
 type blocksStoreSetMock struct {

pkg/querier/blocks_store_queryable_test.go

@@ -7,9 +7,14 @@ import (
 	"testing"
 	"time"
 
+	"github.com/go-kit/log"
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/model/labels"
+	"github.com/prometheus/prometheus/promql"
+	v1 "github.com/prometheus/prometheus/web/api/v1"
 	"github.com/stretchr/testify/require"
+	"github.com/thanos-community/promql-engine/engine"
+	"github.com/thanos-community/promql-engine/logicalplan"
 
 	"github.com/cortexproject/cortex/pkg/chunk"
 	promchunk "github.com/cortexproject/cortex/pkg/chunk/encoding"
@@ -19,14 +24,31 @@ import (
 var _ SeriesWithChunks = &chunkSeries{}
 
 func TestChunkQueryable(t *testing.T) {
-	for _, testcase := range testcases {
-		for _, encoding := range encodings {
-			for _, query := range queries {
-				t.Run(fmt.Sprintf("%s/%s/%s", testcase.name, encoding.name, query.query), func(t *testing.T) {
-					store, from := makeMockChunkStore(t, 24, encoding.e)
-					queryable := newMockStoreQueryable(store, testcase.f)
-					testRangeQuery(t, queryable, from, query)
-				})
+	opts := promql.EngineOpts{
+		Logger:             log.NewNopLogger(),
+		ActiveQueryTracker: promql.NewActiveQueryTracker(t.TempDir(), 10, log.NewNopLogger()),
+		MaxSamples:         1e6,
+		Timeout:            1 * time.Minute,
+	}
+	for _, thanosEngine := range []bool{false, true} {
+		var queryEngine v1.QueryEngine
+		if thanosEngine {
+			queryEngine = engine.New(engine.Opts{
+				EngineOpts:        opts,
+				LogicalOptimizers: logicalplan.AllOptimizers,
+			})
+		} else {
+			queryEngine = promql.NewEngine(opts)
+		}
+		for _, testcase := range testcases {
+			for _, encoding := range encodings {
+				for _, query := range queries {
+					t.Run(fmt.Sprintf("%s/%s/%s/ thanos engine enabled = %t", testcase.name, encoding.name, query.query, thanosEngine), func(t *testing.T) {
+						store, from := makeMockChunkStore(t, 24, encoding.e)
+						queryable := newMockStoreQueryable(store, testcase.f)
+						testRangeQuery(t, queryable, queryEngine, from, query)
+					})
+				}
 			}
 		}
 	}