From babfcfdd910ef2d3e829ae05bac766c1c4e3952b Mon Sep 17 00:00:00 2001
From: Bryan Boreham <bjboreham@gmail.com>
Date: Mon, 29 Apr 2024 13:13:23 +0100
Subject: [PATCH] refactor: Move NewTestEngine into promqltest

And export `DefaultMaxSamplesPerQuery` so callers can replicate previous
behaviour.

Signed-off-by: Bryan Boreham <bjboreham@gmail.com>
---
 promql/engine_test.go          |  2 +-
 promql/promql_test.go          | 12 ++----------
 promql/promqltest/test.go      | 17 ++++++++++++++++-
 promql/promqltest/test_test.go |  2 +-
 4 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/promql/engine_test.go b/promql/engine_test.go
index 36f36807a..1c2722a1d 100644
--- a/promql/engine_test.go
+++ b/promql/engine_test.go
@@ -4999,7 +4999,7 @@ metric 0 1 2
 			if c.engineLookback != 0 {
 				engine.lookbackDelta = c.engineLookback
 			}
-			opts := NewPrometheusQueryOpts(false, c.queryLookback)
+			opts := promql.NewPrometheusQueryOpts(false, c.queryLookback)
 			qry, err := engine.NewInstantQuery(context.Background(), storage, opts, query, c.ts)
 			require.NoError(t, err)
 
diff --git a/promql/promql_test.go b/promql/promql_test.go
index 05821b1c1..87ec2cd4e 100644
--- a/promql/promql_test.go
+++ b/promql/promql_test.go
@@ -22,20 +22,12 @@ import (
 	"github.com/stretchr/testify/require"
 	"golang.org/x/sync/errgroup"
 
+	"github.com/prometheus/prometheus/promql/promqltest"
 	"github.com/prometheus/prometheus/util/teststorage"
 )
 
 func newTestEngine() *Engine {
-	return NewEngine(EngineOpts{
-		Logger:                   nil,
-		Reg:                      nil,
-		MaxSamples:               10000,
-		Timeout:                  100 * time.Second,
-		NoStepSubqueryIntervalFn: func(int64) int64 { return durationMilliseconds(1 * time.Minute) },
-		EnableAtModifier:         true,
-		EnableNegativeOffset:     true,
-		EnablePerStepStats:       true,
-	})
+	return promqltest.NewTestEngine(false, 0, promqltest.DefaultMaxSamplesPerQuery)
 }
 
 func TestEvaluations(t *testing.T) {
diff --git a/promql/promqltest/test.go b/promql/promqltest/test.go
index eab26ad57..1b2ce78af 100644
--- a/promql/promqltest/test.go
+++ b/promql/promqltest/test.go
@@ -49,7 +49,8 @@ var (
 )
 
 const (
-	defaultEpsilon = 0.000001 // Relative error allowed for sample values.
+	defaultEpsilon            = 0.000001 // Relative error allowed for sample values.
+	DefaultMaxSamplesPerQuery = 10000
 )
 
 var testStartTime = time.Unix(0, 0).UTC()
@@ -71,6 +72,20 @@ func LoadedStorage(t testutil.T, input string) *teststorage.TestStorage {
 	return test.storage
 }
 
+func NewTestEngine(enablePerStepStats bool, lookbackDelta time.Duration, maxSamples int) *promql.Engine {
+	return promql.NewEngine(promql.EngineOpts{
+		Logger:                   nil,
+		Reg:                      nil,
+		MaxSamples:               DefaultMaxSamplesPerQuery,
+		Timeout:                  100 * time.Second,
+		NoStepSubqueryIntervalFn: func(int64) int64 { return durationMilliseconds(1 * time.Minute) },
+		EnableAtModifier:         true,
+		EnableNegativeOffset:     true,
+		EnablePerStepStats:       enablePerStepStats,
+		LookbackDelta:            lookbackDelta,
+	})
+}
+
 // RunBuiltinTests runs an acceptance test suite against the provided engine.
 func RunBuiltinTests(t *testing.T, engine promql.QueryEngine) {
 	t.Cleanup(func() { parser.EnableExperimentalFunctions = false })
diff --git a/promql/promqltest/test_test.go b/promql/promqltest/test_test.go
index d28f556f2..f6fe38707 100644
--- a/promql/promqltest/test_test.go
+++ b/promql/promqltest/test_test.go
@@ -451,7 +451,7 @@ eval range from 0 to 5m step 5m testmetric
 
 	for name, testCase := range testCases {
 		t.Run(name, func(t *testing.T) {
-			err := runTest(t, testCase.input, NewTestEngine())
+			err := runTest(t, testCase.input, NewTestEngine(false, 0, DefaultMaxSamplesPerQuery))
 
 			if testCase.expectedError == "" {
 				require.NoError(t, err)