From b0e7de154ee6f6de867fb79739d8406a8bea9b05 Mon Sep 17 00:00:00 2001
From: Achille Roussel <achille.roussel@gmail.com>
Date: Tue, 10 Dec 2024 17:36:44 -0800
Subject: [PATCH] promqltest: allow running tests with custom storage
 implementation

Signed-off-by: Achille Roussel <achille.roussel@gmail.com>
---
 promql/promqltest/test.go | 34 ++++++++++++++++++++++++----------
 1 file changed, 24 insertions(+), 10 deletions(-)

diff --git a/promql/promqltest/test.go b/promql/promqltest/test.go
index 60df005036..efa2136f10 100644
--- a/promql/promqltest/test.go
+++ b/promql/promqltest/test.go
@@ -70,7 +70,7 @@ var testStartTime = time.Unix(0, 0).UTC()
 // LoadedStorage returns storage with generated data using the provided load statements.
 // Non-load statements will cause test errors.
 func LoadedStorage(t testutil.T, input string) *teststorage.TestStorage {
-	test, err := newTest(t, input, false)
+	test, err := newTest(t, input, false, newTestStorage)
 	require.NoError(t, err)
 
 	for _, cmd := range test.cmds {
@@ -81,7 +81,7 @@ func LoadedStorage(t testutil.T, input string) *teststorage.TestStorage {
 			t.Errorf("only 'load' commands accepted, got '%s'", cmd)
 		}
 	}
-	return test.storage
+	return test.storage.(*teststorage.TestStorage)
 }
 
 // NewTestEngine creates a promql.Engine with enablePerStepStats, lookbackDelta and maxSamples, and returns it.
@@ -112,6 +112,11 @@ func NewTestEngineWithOpts(tb testing.TB, opts promql.EngineOpts) *promql.Engine
 
 // RunBuiltinTests runs an acceptance test suite against the provided engine.
 func RunBuiltinTests(t TBRun, engine promql.QueryEngine) {
+	RunBuiltinTestsWithStorage(t, engine, newTestStorage)
+}
+
+// RunBuiltinTestsWithStorage runs an acceptance test suite against the provided engine and storage.
+func RunBuiltinTestsWithStorage(t TBRun, engine promql.QueryEngine, newStorage func(testutil.T) storage.Storage) {
 	t.Cleanup(func() { parser.EnableExperimentalFunctions = false })
 	parser.EnableExperimentalFunctions = true
 
@@ -122,24 +127,29 @@ func RunBuiltinTests(t TBRun, engine promql.QueryEngine) {
 		t.Run(fn, func(t *testing.T) {
 			content, err := fs.ReadFile(testsFs, fn)
 			require.NoError(t, err)
-			RunTest(t, string(content), engine)
+			RunTestWithStorage(t, string(content), engine, newStorage)
 		})
 	}
 }
 
 // RunTest parses and runs the test against the provided engine.
 func RunTest(t testutil.T, input string, engine promql.QueryEngine) {
-	require.NoError(t, runTest(t, input, engine, false))
+	RunTestWithStorage(t, input, engine, newTestStorage)
+}
+
+// RunTestWithStorage parses and runs the test against the provided engine and storage.
+func RunTestWithStorage(t testutil.T, input string, engine promql.QueryEngine, newStorage func(testutil.T) storage.Storage) {
+	require.NoError(t, runTest(t, input, engine, newStorage, false))
 }
 
 // testTest allows tests to be run in "test-the-test" mode (true for
 // testingMode). This is a special mode for testing test code execution itself.
 func testTest(t testutil.T, input string, engine promql.QueryEngine) error {
-	return runTest(t, input, engine, true)
+	return runTest(t, input, engine, newTestStorage, true)
 }
 
-func runTest(t testutil.T, input string, engine promql.QueryEngine, testingMode bool) error {
-	test, err := newTest(t, input, testingMode)
+func runTest(t testutil.T, input string, engine promql.QueryEngine, newStorage func(testutil.T) storage.Storage, testingMode bool) error {
+	test, err := newTest(t, input, testingMode, newStorage)
 
 	// Why do this before checking err? newTest() can create the test storage and then return an error,
 	// and we want to make sure to clean that up to avoid leaking goroutines.
@@ -179,18 +189,20 @@ type test struct {
 
 	cmds []testCommand
 
-	storage *teststorage.TestStorage
+	open    func(testutil.T) storage.Storage
+	storage storage.Storage
 
 	context   context.Context
 	cancelCtx context.CancelFunc
 }
 
 // newTest returns an initialized empty Test.
-func newTest(t testutil.T, input string, testingMode bool) (*test, error) {
+func newTest(t testutil.T, input string, testingMode bool, newStorage func(testutil.T) storage.Storage) (*test, error) {
 	test := &test{
 		T:           t,
 		cmds:        []testCommand{},
 		testingMode: testingMode,
+		open:        newStorage,
 	}
 	err := test.parse(input)
 	test.clear()
@@ -198,6 +210,8 @@ func newTest(t testutil.T, input string, testingMode bool) (*test, error) {
 	return test, err
 }
 
+func newTestStorage(t testutil.T) storage.Storage { return teststorage.New(t) }
+
 //go:embed testdata
 var testsFs embed.FS
 
@@ -1271,7 +1285,7 @@ func (t *test) clear() {
 	if t.cancelCtx != nil {
 		t.cancelCtx()
 	}
-	t.storage = teststorage.New(t)
+	t.storage = t.open(t.T)
 	t.context, t.cancelCtx = context.WithCancel(context.Background())
 }