Fix tests in promql package

This makes tests support the new pushdown interface.

Author: cevian

pkg/promql/bench_test.go

@@ -23,11 +23,10 @@ import (
 
 	"github.com/prometheus/prometheus/pkg/labels"
 	"github.com/prometheus/prometheus/promql/parser"
-	"github.com/prometheus/prometheus/util/teststorage"
 )
 
 func BenchmarkRangeQuery(b *testing.B) {
-	storage := teststorage.New(b)
+	storage := NewTestStorage(b)
 	defer storage.Close()
 	opts := EngineOpts{
 		Logger:     nil,

pkg/promql/engine_test.go

@@ -31,6 +31,13 @@ import (
 	"github.com/prometheus/prometheus/util/testutil"
 )
 
+type QueryableFunc func(ctx context.Context, mint, maxt int64) (Querier, error)
+
+// Querier calls f() with the given parameters.
+func (f QueryableFunc) Querier(ctx context.Context, mint, maxt int64) (Querier, error) {
+	return f(ctx, mint, maxt)
+}
+
 func TestQueryConcurrency(t *testing.T) {
 	maxConcurrency := 10
 
@@ -175,8 +182,8 @@ type errQuerier struct {
 	err error
 }
 
-func (q *errQuerier) Select(bool, *storage.SelectHints, ...*labels.Matcher) (storage.SeriesSet, storage.Warnings, error) {
-	return errSeriesSet{err: q.err}, nil, q.err
+func (q *errQuerier) Select(bool, *storage.SelectHints, []parser.Node, ...*labels.Matcher) (storage.SeriesSet, parser.Node, storage.Warnings, error) {
+	return errSeriesSet{err: q.err}, nil, nil, q.err
 }
 func (*errQuerier) LabelValues(string) ([]string, storage.Warnings, error) { return nil, nil, nil }
 func (*errQuerier) LabelNames() ([]string, storage.Warnings, error)        { return nil, nil, nil }
@@ -200,7 +207,7 @@ func TestQueryError(t *testing.T) {
 	}
 	engine := NewEngine(opts)
 	errStorage := ErrStorage{errors.New("storage error")}
-	queryable := storage.QueryableFunc(func(ctx context.Context, mint, maxt int64) (storage.Querier, error) {
+	queryable := QueryableFunc(func(ctx context.Context, mint, maxt int64) (Querier, error) {
 		return &errQuerier{err: errStorage}, nil
 	})
 	ctx, cancelCtx := context.WithCancel(context.Background())
@@ -234,15 +241,15 @@ type hintCheckerQuerier struct {
 	t *testing.T
 }
 
-func (q *hintCheckerQuerier) Select(_ bool, sp *storage.SelectHints, _ ...*labels.Matcher) (storage.SeriesSet, storage.Warnings, error) {
+func (q *hintCheckerQuerier) Select(_ bool, sp *storage.SelectHints, _ []parser.Node, _ ...*labels.Matcher) (storage.SeriesSet, parser.Node, storage.Warnings, error) {
 	testutil.Equals(q.t, q.start, sp.Start)
 	testutil.Equals(q.t, q.end, sp.End)
 	testutil.Equals(q.t, q.grouping, sp.Grouping)
 	testutil.Equals(q.t, q.by, sp.By)
 	testutil.Equals(q.t, q.selRange, sp.Range)
 	testutil.Equals(q.t, q.function, sp.Func)
 
-	return errSeriesSet{err: nil}, nil, nil
+	return errSeriesSet{err: nil}, nil, nil, nil
 }
 func (*hintCheckerQuerier) LabelValues(string) ([]string, storage.Warnings, error) {
 	return nil, nil, nil
@@ -428,7 +435,7 @@ func TestParamsSetCorrectly(t *testing.T) {
 
 	for _, tc := range cases {
 		engine := NewEngine(opts)
-		queryable := storage.QueryableFunc(func(ctx context.Context, mint, maxt int64) (storage.Querier, error) {
+		queryable := QueryableFunc(func(ctx context.Context, mint, maxt int64) (Querier, error) {
 			return &hintCheckerQuerier{start: tc.paramStart * 1000, end: tc.paramEnd * 1000, grouping: tc.paramGrouping, by: tc.paramBy, selRange: tc.paramRange, function: tc.paramFunc, t: t}, nil
 		})
 

pkg/promql/functions_test.go

@@ -22,15 +22,14 @@ import (
 	"github.com/prometheus/prometheus/pkg/labels"
 	"github.com/prometheus/prometheus/pkg/timestamp"
 	"github.com/prometheus/prometheus/promql/parser"
-	"github.com/prometheus/prometheus/util/teststorage"
 	"github.com/prometheus/prometheus/util/testutil"
 )
 
 func TestDeriv(t *testing.T) {
 	// https://github.com/prometheus/prometheus/issues/2674#issuecomment-315439393
 	// This requires more precision than the usual test system offers,
 	// so we test it by hand.
-	storage := teststorage.New(t)
+	storage := NewTestStorage(t)
 	defer storage.Close()
 	opts := EngineOpts{
 		Logger:     nil,

pkg/promql/test.go

@@ -18,6 +18,7 @@ import (
 	"fmt"
 	"io/ioutil"
 	"math"
+	"os"
 	"regexp"
 	"strconv"
 	"strings"
@@ -47,14 +48,63 @@ const (
 
 var testStartTime = time.Unix(0, 0).UTC()
 
+// New returns a new TestStorage for testing purposes
+// that removes all associated files on closing.
+func NewTestStorage(t testutil.T) *TestStorage {
+	dir, err := ioutil.TempDir("", "test_storage")
+	if err != nil {
+		t.Fatalf("Opening test dir failed: %s", err)
+	}
+
+	// Tests just load data for a series sequentially. Thus we
+	// need a long appendable window.
+	opts := tsdb.DefaultOptions()
+	opts.MinBlockDuration = int64(24 * time.Hour / time.Millisecond)
+	opts.MaxBlockDuration = int64(24 * time.Hour / time.Millisecond)
+	db, err := tsdb.Open(dir, nil, nil, opts)
+	if err != nil {
+		t.Fatalf("Opening test storage failed: %s", err)
+	}
+	return &TestStorage{DB: db, dir: dir}
+}
+
+type TestStorage struct {
+	*tsdb.DB
+	dir string
+}
+
+type QuerierWrapper struct {
+	storage.Querier
+}
+
+func (t *QuerierWrapper) Select(b bool, sh *storage.SelectHints, _ []parser.Node, m ...*labels.Matcher) (storage.SeriesSet, parser.Node, storage.Warnings, error) {
+	ss, w, err := t.Querier.Select(b, sh, m...)
+	return ss, nil, w, err
+}
+
+func (db *TestStorage) Querier(ctx context.Context, mint, maxt int64) (Querier, error) {
+	q, err := db.DB.Querier(ctx, mint, maxt)
+	if err != nil {
+		return nil, err
+	}
+	return &QuerierWrapper{q}, err
+}
+
+func (s TestStorage) Close() error {
+	if err := s.DB.Close(); err != nil {
+		return err
+	}
+	return os.RemoveAll(s.dir)
+}
+
 // Test is a sequence of read and write commands that are run
 // against a test storage.
 type Test struct {
 	testutil.T
 
 	cmds []testCommand
 
-	storage *teststorage.TestStorage
+	storage *TestStorage
 
 	queryEngine *Engine
 	context     context.Context
@@ -87,7 +137,7 @@ func (t *Test) QueryEngine() *Engine {
 }
 
 // Queryable allows querying the test data.
-func (t *Test) Queryable() storage.Queryable {
+func (t *Test) Queryable() Queryable {
 	return t.storage
 }
 
@@ -98,7 +148,7 @@ func (t *Test) Context() context.Context {
 
 // Storage returns the test's storage.
 func (t *Test) Storage() storage.Storage {
-	return t.storage
+	return t.storage.DB
 }
 
 // TSDB returns test's TSDB.
@@ -515,7 +565,7 @@ func (t *Test) clear() {
 	if t.cancelCtx != nil {
 		t.cancelCtx()
 	}
-	t.storage = teststorage.New(t)
+	t.storage = NewTestStorage(t)
 
 	opts := EngineOpts{
 		Logger:     nil,