Refactor engine creation in tests

Signed-off-by: Arve Knudsen <arve.knudsen@gmail.com>

promql/bench_test.go

@@ -21,12 +21,11 @@ import (
 	"testing"
 	"time"
 
-	"github.com/stretchr/testify/require"
-
 	"github.com/prometheus/prometheus/model/histogram"
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/promql"
 	"github.com/prometheus/prometheus/promql/parser"
+	"github.com/prometheus/prometheus/promql/promqltest"
 	"github.com/prometheus/prometheus/storage"
 	"github.com/prometheus/prometheus/tsdb/tsdbutil"
 	"github.com/prometheus/prometheus/util/teststorage"
@@ -273,10 +272,7 @@ func BenchmarkRangeQuery(b *testing.B) {
 		MaxSamples: 50000000,
 		Timeout:    100 * time.Second,
 	}
-	engine := promql.NewEngine(opts)
+	engine := promqltest.NewTestEngineWithOpts(b, opts)
-	b.Cleanup(func() {
-		require.NoError(b, engine.Close())
-	})
 
 	const interval = 10000 // 10s interval.
 	// A day of data plus 10k steps.
@@ -367,10 +363,7 @@ func BenchmarkNativeHistograms(b *testing.B) {
 
 	for _, tc := range cases {
 		b.Run(tc.name, func(b *testing.B) {
-			ng := promql.NewEngine(opts)
+			ng := promqltest.NewTestEngineWithOpts(b, opts)
-			b.Cleanup(func() {
-				require.NoError(b, ng.Close())
-			})
 			for i := 0; i < b.N; i++ {
 				qry, err := ng.NewRangeQuery(context.Background(), testStorage, nil, tc.query, start, end, step)
 				if err != nil {

promql/engine_test.go

@@ -64,11 +64,8 @@ func TestQueryConcurrency(t *testing.T) {
 		Timeout:            100 * time.Second,
 		ActiveQueryTracker: queryTracker,
 	}
+	engine := promqltest.NewTestEngineWithOpts(t, opts)
 
-	engine := promql.NewEngine(opts)
-	t.Cleanup(func() {
-		require.NoError(t, engine.Close())
-	})
 	ctx, cancelCtx := context.WithCancel(context.Background())
 	t.Cleanup(cancelCtx)
 
@@ -162,10 +159,7 @@ func TestQueryTimeout(t *testing.T) {
 		MaxSamples: 10,
 		Timeout:    5 * time.Millisecond,
 	}
-	engine := promql.NewEngine(opts)
+	engine := promqltest.NewTestEngineWithOpts(t, opts)
-	t.Cleanup(func() {
-		require.NoError(t, engine.Close())
-	})
 	ctx, cancelCtx := context.WithCancel(context.Background())
 	defer cancelCtx()
 
@@ -190,7 +184,7 @@ func TestQueryCancel(t *testing.T) {
 		MaxSamples: 10,
 		Timeout:    10 * time.Second,
 	}
-	engine := promql.NewEngine(opts)
+	engine := promqltest.NewTestEngineWithOpts(t, opts)
 	ctx, cancelCtx := context.WithCancel(context.Background())
 	defer cancelCtx()
 
@@ -264,10 +258,7 @@ func TestQueryError(t *testing.T) {
 		MaxSamples: 10,
 		Timeout:    10 * time.Second,
 	}
-	engine := promql.NewEngine(opts)
+	engine := promqltest.NewTestEngineWithOpts(t, opts)
-	t.Cleanup(func() {
-		require.NoError(t, engine.Close())
-	})
 	errStorage := promql.ErrStorage{errors.New("storage error")}
 	queryable := storage.QueryableFunc(func(mint, maxt int64) (storage.Querier, error) {
 		return &errQuerier{err: errStorage}, nil
@@ -601,10 +592,7 @@ func TestSelectHintsSetCorrectly(t *testing.T) {
 		},
 	} {
 		t.Run(tc.query, func(t *testing.T) {
-			engine := promql.NewEngine(opts)
+			engine := promqltest.NewTestEngineWithOpts(t, opts)
-			t.Cleanup(func() {
-				require.NoError(t, engine.Close())
-			})
 			hintsRecorder := &noopHintRecordingQueryable{}
 
 			var (
@@ -635,10 +623,7 @@ func TestEngineShutdown(t *testing.T) {
 		MaxSamples: 10,
 		Timeout:    10 * time.Second,
 	}
-	engine := promql.NewEngine(opts)
+	engine := promqltest.NewTestEngineWithOpts(t, opts)
-	t.Cleanup(func() {
-		require.NoError(t, engine.Close())
-	})
 	ctx, cancelCtx := context.WithCancel(context.Background())
 
 	block := make(chan struct{})
@@ -2055,10 +2040,7 @@ func TestQueryLogger_basic(t *testing.T) {
 		MaxSamples: 10,
 		Timeout:    10 * time.Second,
 	}
-	engine := promql.NewEngine(opts)
+	engine := promqltest.NewTestEngineWithOpts(t, opts)
-	t.Cleanup(func() {
-		require.NoError(t, engine.Close())
-	})
 
 	queryExec := func() {
 		ctx, cancelCtx := context.WithCancel(context.Background())
@@ -2109,10 +2091,7 @@ func TestQueryLogger_fields(t *testing.T) {
 		MaxSamples: 10,
 		Timeout:    10 * time.Second,
 	}
-	engine := promql.NewEngine(opts)
+	engine := promqltest.NewTestEngineWithOpts(t, opts)
-	t.Cleanup(func() {
-		require.NoError(t, engine.Close())
-	})
 
 	f1 := NewFakeQueryLogger()
 	engine.SetQueryLogger(f1)
@@ -2141,10 +2120,7 @@ func TestQueryLogger_error(t *testing.T) {
 		MaxSamples: 10,
 		Timeout:    10 * time.Second,
 	}
-	engine := promql.NewEngine(opts)
+	engine := promqltest.NewTestEngineWithOpts(t, opts)
-	t.Cleanup(func() {
-		require.NoError(t, engine.Close())
-	})
 
 	f1 := NewFakeQueryLogger()
 	engine.SetQueryLogger(f1)
@@ -3027,10 +3003,7 @@ func TestEngineOptsValidation(t *testing.T) {
 	}
 
 	for _, c := range cases {
-		eng := promql.NewEngine(c.opts)
+		eng := promqltest.NewTestEngineWithOpts(t, c.opts)
-		t.Cleanup(func() {
-			require.NoError(t, eng.Close())
-		})
 		_, err1 := eng.NewInstantQuery(context.Background(), nil, nil, c.query, time.Unix(10, 0))
 		_, err2 := eng.NewRangeQuery(context.Background(), nil, nil, c.query, time.Unix(0, 0), time.Unix(10, 0), time.Second)
 		if c.fail {

promql/functions_test.go

@@ -24,6 +24,7 @@ import (
 	"github.com/prometheus/prometheus/model/timestamp"
 	"github.com/prometheus/prometheus/promql"
 	"github.com/prometheus/prometheus/promql/parser"
+	"github.com/prometheus/prometheus/promql/promqltest"
 	"github.com/prometheus/prometheus/util/teststorage"
 )
 
@@ -39,10 +40,7 @@ func TestDeriv(t *testing.T) {
 		MaxSamples: 10000,
 		Timeout:    10 * time.Second,
 	}
-	engine := promql.NewEngine(opts)
+	engine := promqltest.NewTestEngineWithOpts(t, opts)
-	t.Cleanup(func() {
-		require.NoError(t, engine.Close())
-	})
 
 	a := storage.Appender(context.Background())
 

promql/promql_test.go

@@ -48,10 +48,7 @@ func TestConcurrentRangeQueries(t *testing.T) {
 	}
 	// Enable experimental functions testing
 	parser.EnableExperimentalFunctions = true
-	engine := promql.NewEngine(opts)
+	engine := promqltest.NewTestEngineWithOpts(t, opts)
-	t.Cleanup(func() {
-		require.NoError(t, engine.Close())
-	})
 
 	const interval = 10000 // 10s interval.
 	// A day of data plus 10k steps.

promql/promqltest/test.go

@@ -74,8 +74,9 @@ func LoadedStorage(t testutil.T, input string) *teststorage.TestStorage {
 	return test.storage
 }
 
-func NewTestEngine(t *testing.T, enablePerStepStats bool, lookbackDelta time.Duration, maxSamples int) *promql.Engine {
+// NewTestEngine creates a promql.Engine with enablePerStepStats, lookbackDelta and maxSamples, and returns it.
-	ng := promql.NewEngine(promql.EngineOpts{
+func NewTestEngine(tb testing.TB, enablePerStepStats bool, lookbackDelta time.Duration, maxSamples int) *promql.Engine {
+	return NewTestEngineWithOpts(tb, promql.EngineOpts{
 		Logger:                   nil,
 		Reg:                      nil,
 		MaxSamples:               maxSamples,
@@ -86,8 +87,14 @@ func NewTestEngine(t *testing.T, enablePerStepStats bool, lookbackDelta time.Dur
 		EnablePerStepStats:       enablePerStepStats,
 		LookbackDelta:            lookbackDelta,
 	})
-	t.Cleanup(func() {
+}
-		require.NoError(t, ng.Close())
+
+// NewTestEngineWithOpts creates a promql.Engine with opts and returns it.
+func NewTestEngineWithOpts(tb testing.TB, opts promql.EngineOpts) *promql.Engine {
+	tb.Helper()
+	ng := promql.NewEngine(opts)
+	tb.Cleanup(func() {
+		require.NoError(tb, ng.Close())
 	})
 	return ng
 }

rules/alerting_test.go

@@ -38,7 +38,7 @@ import (
 
 func testEngine(tb testing.TB) *promql.Engine {
 	tb.Helper()
-	e := promql.NewEngine(promql.EngineOpts{
+	return promqltest.NewTestEngineWithOpts(tb, promql.EngineOpts{
 		Logger:                   nil,
 		Reg:                      nil,
 		MaxSamples:               10000,
@@ -48,10 +48,6 @@ func testEngine(tb testing.TB) *promql.Engine {
 		EnableNegativeOffset:     true,
 		EnablePerStepStats:       true,
 	})
-	tb.Cleanup(func() {
-		require.NoError(tb, e.Close())
-	})
-	return e
 }
 
 func TestAlertingRuleState(t *testing.T) {
@@ -595,10 +591,7 @@ func TestAlertingRuleDuplicate(t *testing.T) {
 		Timeout:    10 * time.Second,
 	}
 
-	engine := promql.NewEngine(opts)
+	engine := promqltest.NewTestEngineWithOpts(t, opts)
-	t.Cleanup(func() {
-		require.NoError(t, engine.Close())
-	})
 	ctx, cancelCtx := context.WithCancel(context.Background())
 	defer cancelCtx()
 

rules/manager_test.go

@@ -542,10 +542,7 @@ func TestStaleness(t *testing.T) {
 			MaxSamples: 10,
 			Timeout:    10 * time.Second,
 		}
-		engine := promql.NewEngine(engineOpts)
+		engine := promqltest.NewTestEngineWithOpts(t, engineOpts)
-		t.Cleanup(func() {
-			require.NoError(t, engine.Close())
-		})
 		opts := &ManagerOptions{
 			QueryFunc:  EngineQueryFunc(engine, st),
 			Appendable: st,
@@ -779,10 +776,7 @@ func TestUpdate(t *testing.T) {
 		MaxSamples: 10,
 		Timeout:    10 * time.Second,
 	}
-	engine := promql.NewEngine(opts)
+	engine := promqltest.NewTestEngineWithOpts(t, opts)
-	t.Cleanup(func() {
-		require.NoError(t, engine.Close())
-	})
 	ruleManager := NewManager(&ManagerOptions{
 		Appendable: st,
 		Queryable:  st,
@@ -920,10 +914,7 @@ func TestNotify(t *testing.T) {
 		MaxSamples: 10,
 		Timeout:    10 * time.Second,
 	}
-	engine := promql.NewEngine(engineOpts)
+	engine := promqltest.NewTestEngineWithOpts(t, engineOpts)
-	t.Cleanup(func() {
-		require.NoError(t, engine.Close())
-	})
 	var lastNotified []*Alert
 	notifyFunc := func(ctx context.Context, expr string, alerts ...*Alert) {
 		lastNotified = alerts
@@ -998,10 +989,7 @@ func TestMetricsUpdate(t *testing.T) {
 		MaxSamples: 10,
 		Timeout:    10 * time.Second,
 	}
-	engine := promql.NewEngine(opts)
+	engine := promqltest.NewTestEngineWithOpts(t, opts)
-	t.Cleanup(func() {
-		require.NoError(t, engine.Close())
-	})
 	ruleManager := NewManager(&ManagerOptions{
 		Appendable: storage,
 		Queryable:  storage,
@@ -1075,10 +1063,7 @@ func TestGroupStalenessOnRemoval(t *testing.T) {
 		MaxSamples: 10,
 		Timeout:    10 * time.Second,
 	}
-	engine := promql.NewEngine(opts)
+	engine := promqltest.NewTestEngineWithOpts(t, opts)
-	t.Cleanup(func() {
-		require.NoError(t, engine.Close())
-	})
 	ruleManager := NewManager(&ManagerOptions{
 		Appendable: storage,
 		Queryable:  storage,
@@ -1155,10 +1140,7 @@ func TestMetricsStalenessOnManagerShutdown(t *testing.T) {
 		MaxSamples: 10,
 		Timeout:    10 * time.Second,
 	}
-	engine := promql.NewEngine(opts)
+	engine := promqltest.NewTestEngineWithOpts(t, opts)
-	t.Cleanup(func() {
-		require.NoError(t, engine.Close())
-	})
 	ruleManager := NewManager(&ManagerOptions{
 		Appendable: storage,
 		Queryable:  storage,
@@ -1260,10 +1242,7 @@ func TestRuleHealthUpdates(t *testing.T) {
 		MaxSamples: 10,
 		Timeout:    10 * time.Second,
 	}
-	engine := promql.NewEngine(engineOpts)
+	engine := promqltest.NewTestEngineWithOpts(t, engineOpts)
-	t.Cleanup(func() {
-		require.NoError(t, engine.Close())
-	})
 	opts := &ManagerOptions{
 		QueryFunc:  EngineQueryFunc(engine, st),
 		Appendable: st,

rules/recording_test.go

@@ -167,10 +167,7 @@ func TestRuleEvalDuplicate(t *testing.T) {
 		Timeout:    10 * time.Second,
 	}
 
-	engine := promql.NewEngine(opts)
+	engine := promqltest.NewTestEngineWithOpts(t, opts)
-	t.Cleanup(func() {
-		require.NoError(t, engine.Close())
-	})
 	ctx, cancelCtx := context.WithCancel(context.Background())
 	defer cancelCtx()
 

web/api/v1/api_test.go

@@ -61,8 +61,7 @@ import (
 
 func testEngine(t *testing.T) *promql.Engine {
 	t.Helper()
-
+	return promqltest.NewTestEngineWithOpts(t, promql.EngineOpts{
-	ng := promql.NewEngine(promql.EngineOpts{
 		Logger:                   nil,
 		Reg:                      nil,
 		MaxSamples:               10000,
@@ -72,11 +71,6 @@ func testEngine(t *testing.T) *promql.Engine {
 		EnableNegativeOffset:     true,
 		EnablePerStepStats:       true,
 	})
-	t.Cleanup(func() {
-		require.NoError(t, ng.Close())
-	})
-
-	return ng
 }
 
 // testMetaStore satisfies the scrape.MetricMetadataStore interface.
@@ -315,11 +309,7 @@ func (m *rulesRetrieverMock) CreateRuleGroups() {
 		MaxSamples: 10,
 		Timeout:    100 * time.Second,
 	}
-
+	engine := promqltest.NewTestEngineWithOpts(m.testing, engineOpts)
-	engine := promql.NewEngine(engineOpts)
-	m.testing.Cleanup(func() {
-		require.NoError(m.testing, engine.Close())
-	})
 	opts := &rules.ManagerOptions{
 		QueryFunc:  rules.EngineQueryFunc(engine, storage),
 		Appendable: storage,

web/api/v1/errors_test.go

@@ -32,6 +32,7 @@ import (
 	"github.com/prometheus/prometheus/config"
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/promql"
+	"github.com/prometheus/prometheus/promql/promqltest"
 	"github.com/prometheus/prometheus/rules"
 	"github.com/prometheus/prometheus/scrape"
 	"github.com/prometheus/prometheus/storage"
@@ -103,16 +104,13 @@ func TestApiStatusCodes(t *testing.T) {
 func createPrometheusAPI(t *testing.T, q storage.SampleAndChunkQueryable) *route.Router {
 	t.Helper()
 
-	engine := promql.NewEngine(promql.EngineOpts{
+	engine := promqltest.NewTestEngineWithOpts(t, promql.EngineOpts{
 		Logger:             log.NewNopLogger(),
 		Reg:                nil,
 		ActiveQueryTracker: nil,
 		MaxSamples:         100,
 		Timeout:            5 * time.Second,
 	})
-	t.Cleanup(func() {
-		require.NoError(t, engine.Close())
-	})
 
 	api := NewAPI(
 		engine,