From fbc9eddfaf9ea5136c1008d997834cb3ac07ba90 Mon Sep 17 00:00:00 2001 From: Arve Knudsen Date: Sun, 14 Jul 2024 13:28:59 +0200 Subject: [PATCH] Refactor engine creation in tests Signed-off-by: Arve Knudsen --- promql/bench_test.go | 13 +++-------- promql/engine_test.go | 47 +++++++++------------------------------ promql/functions_test.go | 6 ++--- promql/promql_test.go | 5 +---- promql/promqltest/test.go | 15 +++++++++---- rules/alerting_test.go | 11 ++------- rules/manager_test.go | 35 ++++++----------------------- rules/recording_test.go | 5 +---- web/api/v1/api_test.go | 14 ++---------- web/api/v1/errors_test.go | 6 ++--- 10 files changed, 41 insertions(+), 116 deletions(-) diff --git a/promql/bench_test.go b/promql/bench_test.go index 2f244be1d..bc953464e 100644 --- a/promql/bench_test.go +++ b/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) - b.Cleanup(func() { - require.NoError(b, engine.Close()) - }) + engine := promqltest.NewTestEngineWithOpts(b, opts) 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) - b.Cleanup(func() { - require.NoError(b, ng.Close()) - }) + ng := promqltest.NewTestEngineWithOpts(b, opts) for i := 0; i < b.N; i++ { qry, err := ng.NewRangeQuery(context.Background(), testStorage, nil, tc.query, start, end, step) if err != nil { diff --git a/promql/engine_test.go b/promql/engine_test.go index f166b6d13..d58d23d45 100644 --- a/promql/engine_test.go +++ b/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) - t.Cleanup(func() { - require.NoError(t, engine.Close()) - }) + engine := promqltest.NewTestEngineWithOpts(t, opts) 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) - t.Cleanup(func() { - require.NoError(t, engine.Close()) - }) + engine := promqltest.NewTestEngineWithOpts(t, opts) 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) - t.Cleanup(func() { - require.NoError(t, engine.Close()) - }) + engine := promqltest.NewTestEngineWithOpts(t, opts) hintsRecorder := &noopHintRecordingQueryable{} var ( @@ -635,10 +623,7 @@ func TestEngineShutdown(t *testing.T) { MaxSamples: 10, Timeout: 10 * time.Second, } - engine := promql.NewEngine(opts) - t.Cleanup(func() { - require.NoError(t, engine.Close()) - }) + engine := promqltest.NewTestEngineWithOpts(t, opts) 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) - t.Cleanup(func() { - require.NoError(t, engine.Close()) - }) + engine := promqltest.NewTestEngineWithOpts(t, opts) 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) - t.Cleanup(func() { - require.NoError(t, engine.Close()) - }) + engine := promqltest.NewTestEngineWithOpts(t, opts) 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) - t.Cleanup(func() { - require.NoError(t, engine.Close()) - }) + engine := promqltest.NewTestEngineWithOpts(t, opts) f1 := NewFakeQueryLogger() engine.SetQueryLogger(f1) @@ -3027,10 +3003,7 @@ func TestEngineOptsValidation(t *testing.T) { } for _, c := range cases { - eng := promql.NewEngine(c.opts) - t.Cleanup(func() { - require.NoError(t, eng.Close()) - }) + eng := promqltest.NewTestEngineWithOpts(t, c.opts) _, 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 { diff --git a/promql/functions_test.go b/promql/functions_test.go index ce62d316b..9ee0ba51d 100644 --- a/promql/functions_test.go +++ b/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) - t.Cleanup(func() { - require.NoError(t, engine.Close()) - }) + engine := promqltest.NewTestEngineWithOpts(t, opts) a := storage.Appender(context.Background()) diff --git a/promql/promql_test.go b/promql/promql_test.go index 2a5834681..345ecab5e 100644 --- a/promql/promql_test.go +++ b/promql/promql_test.go @@ -48,10 +48,7 @@ func TestConcurrentRangeQueries(t *testing.T) { } // Enable experimental functions testing parser.EnableExperimentalFunctions = true - engine := promql.NewEngine(opts) - t.Cleanup(func() { - require.NoError(t, engine.Close()) - }) + engine := promqltest.NewTestEngineWithOpts(t, opts) const interval = 10000 // 10s interval. // A day of data plus 10k steps. diff --git a/promql/promqltest/test.go b/promql/promqltest/test.go index 1fda9985a..8c7a346c2 100644 --- a/promql/promqltest/test.go +++ b/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 { - ng := promql.NewEngine(promql.EngineOpts{ +// NewTestEngine creates a promql.Engine with enablePerStepStats, lookbackDelta and maxSamples, and returns it. +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 } diff --git a/rules/alerting_test.go b/rules/alerting_test.go index 2ff641b4b..81c065eaa 100644 --- a/rules/alerting_test.go +++ b/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) - t.Cleanup(func() { - require.NoError(t, engine.Close()) - }) + engine := promqltest.NewTestEngineWithOpts(t, opts) ctx, cancelCtx := context.WithCancel(context.Background()) defer cancelCtx() diff --git a/rules/manager_test.go b/rules/manager_test.go index 854048281..ea65a943d 100644 --- a/rules/manager_test.go +++ b/rules/manager_test.go @@ -542,10 +542,7 @@ func TestStaleness(t *testing.T) { MaxSamples: 10, Timeout: 10 * time.Second, } - engine := promql.NewEngine(engineOpts) - t.Cleanup(func() { - require.NoError(t, engine.Close()) - }) + engine := promqltest.NewTestEngineWithOpts(t, engineOpts) 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) - t.Cleanup(func() { - require.NoError(t, engine.Close()) - }) + engine := promqltest.NewTestEngineWithOpts(t, opts) 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) - t.Cleanup(func() { - require.NoError(t, engine.Close()) - }) + engine := promqltest.NewTestEngineWithOpts(t, engineOpts) 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) - t.Cleanup(func() { - require.NoError(t, engine.Close()) - }) + engine := promqltest.NewTestEngineWithOpts(t, opts) 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) - t.Cleanup(func() { - require.NoError(t, engine.Close()) - }) + engine := promqltest.NewTestEngineWithOpts(t, opts) 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) - t.Cleanup(func() { - require.NoError(t, engine.Close()) - }) + engine := promqltest.NewTestEngineWithOpts(t, opts) 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) - t.Cleanup(func() { - require.NoError(t, engine.Close()) - }) + engine := promqltest.NewTestEngineWithOpts(t, engineOpts) opts := &ManagerOptions{ QueryFunc: EngineQueryFunc(engine, st), Appendable: st, diff --git a/rules/recording_test.go b/rules/recording_test.go index 98640ab0e..72c0764f9 100644 --- a/rules/recording_test.go +++ b/rules/recording_test.go @@ -167,10 +167,7 @@ func TestRuleEvalDuplicate(t *testing.T) { Timeout: 10 * time.Second, } - engine := promql.NewEngine(opts) - t.Cleanup(func() { - require.NoError(t, engine.Close()) - }) + engine := promqltest.NewTestEngineWithOpts(t, opts) ctx, cancelCtx := context.WithCancel(context.Background()) defer cancelCtx() diff --git a/web/api/v1/api_test.go b/web/api/v1/api_test.go index 3e94d18a0..61617189d 100644 --- a/web/api/v1/api_test.go +++ b/web/api/v1/api_test.go @@ -61,8 +61,7 @@ import ( func testEngine(t *testing.T) *promql.Engine { t.Helper() - - ng := promql.NewEngine(promql.EngineOpts{ + return promqltest.NewTestEngineWithOpts(t, 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 := promql.NewEngine(engineOpts) - m.testing.Cleanup(func() { - require.NoError(m.testing, engine.Close()) - }) + engine := promqltest.NewTestEngineWithOpts(m.testing, engineOpts) opts := &rules.ManagerOptions{ QueryFunc: rules.EngineQueryFunc(engine, storage), Appendable: storage, diff --git a/web/api/v1/errors_test.go b/web/api/v1/errors_test.go index d7fee4079..7e1fc09d8 100644 --- a/web/api/v1/errors_test.go +++ b/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,