promqltest: allow running tests with custom storage implementation

Signed-off-by: Achille Roussel <achille.roussel@gmail.com>
This commit is contained in:
Achille Roussel 2024-12-10 17:36:44 -08:00
parent e1bfdd2257
commit b0e7de154e

View File

@ -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())
}