Promtool: Add feature flags to unit tests (#8958)

* Added feature flag support to unit tests

Signed-off-by: Levi Harrison <git@leviharrison.dev>

* Added/fixed tests

Signed-off-by: Levi Harrison <git@leviharrison.dev>

* Addressed review comments

Signed-off-by: Levi Harrison <git@leviharrison.dev>

cmd/promtool/main.go

@@ -49,6 +49,7 @@ import (
 	"github.com/prometheus/prometheus/discovery/kubernetes"
 	"github.com/prometheus/prometheus/discovery/targetgroup"
 	"github.com/prometheus/prometheus/pkg/rulefmt"
+	"github.com/prometheus/prometheus/promql"
 )
 
 func main() {
@@ -165,6 +166,8 @@ func main() {
 		"A list of one or more files containing recording rules to be backfilled. All recording rules listed in the files will be backfilled. Alerting rules are not evaluated.",
 	).Required().ExistingFiles()
 
+	featureList := app.Flag("enable-feature", "Comma separated feature names to enable (only PromQL related). See https://prometheus.io/docs/prometheus/latest/disabled_features/ for the options and more details.").Default("").Strings()
+
 	parsedCmd := kingpin.MustParse(app.Parse(os.Args[1:]))
 
 	var p printer
@@ -175,6 +178,23 @@ func main() {
 		p = &promqlPrinter{}
 	}
 
+	var queryOpts promql.LazyLoaderOpts
+	for _, f := range *featureList {
+		opts := strings.Split(f, ",")
+		for _, o := range opts {
+			switch o {
+			case "promql-at-modifier":
+				queryOpts.EnableAtModifier = true
+			case "promql-negative-offset":
+				queryOpts.EnableNegativeOffset = true
+			case "":
+				continue
+			default:
+				fmt.Printf("  WARNING: Unknown option for --enable-feature: %q\n", o)
+			}
+		}
+	}
+
 	switch parsedCmd {
 	case checkConfigCmd.FullCommand():
 		os.Exit(CheckConfig(*configFiles...))
@@ -210,7 +230,7 @@ func main() {
 		os.Exit(QueryLabels(*queryLabelsServer, *queryLabelsName, *queryLabelsBegin, *queryLabelsEnd, p))
 
 	case testRulesCmd.FullCommand():
-		os.Exit(RulesUnitTest(*testRulesFiles...))
+		os.Exit(RulesUnitTest(queryOpts, *testRulesFiles...))
 
 	case tsdbBenchWriteCmd.FullCommand():
 		os.Exit(checkErr(benchmarkWrite(*benchWriteOutPath, *benchSamplesFile, *benchWriteNumMetrics, *benchWriteNumScrapes)))

cmd/promtool/testdata/at-modifier-test.yml

@@ -0,0 +1,7 @@
+rule_files:
+  - at-modifier.yml
+
+tests:
+  - input_series:
+      - series: "requests{}"
+        values: 1

cmd/promtool/testdata/at-modifier.yml

@@ -0,0 +1,7 @@
+# This is the rules file for at-modifier-test.yml.
+
+groups:
+  - name: at-modifier
+    rules:
+      - record: x
+        expr: "requests @ 1000"

cmd/promtool/testdata/negative-offset-test.yml

@@ -0,0 +1,7 @@
+rule_files:
+  - negative-offset.yml
+
+tests:
+  - input_series:
+      - series: "requests{}"
+        values: 1

cmd/promtool/testdata/negative-offset.yml

@@ -0,0 +1,7 @@
+# This is the rules file for negative-offset-test.yml.
+
+groups:
+  - name: negative-offset
+    rules:
+      - record: x
+        expr: "requests offset -5m"

cmd/promtool/unittest.go

@@ -39,11 +39,11 @@ import (
 
 // RulesUnitTest does unit testing of rules based on the unit testing files provided.
 // More info about the file format can be found in the docs.
-func RulesUnitTest(files ...string) int {
+func RulesUnitTest(queryOpts promql.LazyLoaderOpts, files ...string) int {
 	failed := false
 
 	for _, f := range files {
-		if errs := ruleUnitTest(f); errs != nil {
+		if errs := ruleUnitTest(f, queryOpts); errs != nil {
 			fmt.Fprintln(os.Stderr, "  FAILED:")
 			for _, e := range errs {
 				fmt.Fprintln(os.Stderr, e.Error())
@@ -60,7 +60,7 @@ func RulesUnitTest(files ...string) int {
 	return 0
 }
 
-func ruleUnitTest(filename string) []error {
+func ruleUnitTest(filename string, queryOpts promql.LazyLoaderOpts) []error {
 	fmt.Println("Unit Testing: ", filename)
 
 	b, err := ioutil.ReadFile(filename)
@@ -95,7 +95,7 @@ func ruleUnitTest(filename string) []error {
 	// Testing.
 	var errs []error
 	for _, t := range unitTestInp.Tests {
-		ers := t.test(evalInterval, groupOrderMap, unitTestInp.RuleFiles...)
+		ers := t.test(evalInterval, groupOrderMap, queryOpts, unitTestInp.RuleFiles...)
 		if ers != nil {
 			errs = append(errs, ers...)
 		}
@@ -151,9 +151,9 @@ type testGroup struct {
 }
 
 // test performs the unit tests.
-func (tg *testGroup) test(evalInterval time.Duration, groupOrderMap map[string]int, ruleFiles ...string) []error {
+func (tg *testGroup) test(evalInterval time.Duration, groupOrderMap map[string]int, queryOpts promql.LazyLoaderOpts, ruleFiles ...string) []error {
 	// Setup testing suite.
-	suite, err := promql.NewLazyLoader(nil, tg.seriesLoadingString())
+	suite, err := promql.NewLazyLoader(nil, tg.seriesLoadingString(), queryOpts)
 	if err != nil {
 		return []error{err}
 	}

cmd/promtool/unittest_test.go

@@ -13,7 +13,11 @@
 
 package main
 
-import "testing"
+import (
+	"testing"
+
+	"github.com/prometheus/prometheus/promql"
+)
 
 func TestRulesUnitTest(t *testing.T) {
 	type args struct {
@@ -22,6 +26,7 @@ func TestRulesUnitTest(t *testing.T) {
 	tests := []struct {
 		name      string
 		args      args
+		queryOpts promql.LazyLoaderOpts
 		want      int
 	}{
 		{
@@ -66,10 +71,44 @@ func TestRulesUnitTest(t *testing.T) {
 			},
 			want: 1,
 		},
+		{
+			name: "Disabled feature (@ modifier)",
+			args: args{
+				files: []string{"./testdata/at-modifier-test.yml"},
+			},
+			want: 1,
+		},
+		{
+			name: "Enabled feature (@ modifier)",
+			args: args{
+				files: []string{"./testdata/at-modifier-test.yml"},
+			},
+			queryOpts: promql.LazyLoaderOpts{
+				EnableAtModifier: true,
+			},
+			want: 0,
+		},
+		{
+			name: "Disabled feature (negative offset)",
+			args: args{
+				files: []string{"./testdata/negative-offset-test.yml"},
+			},
+			want: 1,
+		},
+		{
+			name: "Enabled feature (negative offset)",
+			args: args{
+				files: []string{"./testdata/negative-offset-test.yml"},
+			},
+			queryOpts: promql.LazyLoaderOpts{
+				EnableNegativeOffset: true,
+			},
+			want: 0,
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			if got := RulesUnitTest(tt.args.files...); got != tt.want {
+			if got := RulesUnitTest(tt.queryOpts, tt.args.files...); got != tt.want {
 				t.Errorf("RulesUnitTest() = %v, want %v", got, tt.want)
 			}
 		})

promql/test.go

@@ -675,12 +675,21 @@ type LazyLoader struct {
 	queryEngine *Engine
 	context     context.Context
 	cancelCtx   context.CancelFunc
+
+	opts LazyLoaderOpts
+}
+
+// LazyLoaderOpts are options for the lazy loader.
+type LazyLoaderOpts struct {
+	// Disabled PromQL engine features.
+	EnableAtModifier, EnableNegativeOffset bool
 }
 
 // NewLazyLoader returns an initialized empty LazyLoader.
-func NewLazyLoader(t testutil.T, input string) (*LazyLoader, error) {
+func NewLazyLoader(t testutil.T, input string, opts LazyLoaderOpts) (*LazyLoader, error) {
 	ll := &LazyLoader{
 		T:    t,
+		opts: opts,
 	}
 	err := ll.parse(input)
 	ll.clear()
@@ -728,7 +737,8 @@ func (ll *LazyLoader) clear() {
 		MaxSamples:               10000,
 		Timeout:                  100 * time.Second,
 		NoStepSubqueryIntervalFn: func(int64) int64 { return durationMilliseconds(ll.SubqueryInterval) },
-		EnableAtModifier:         true,
+		EnableAtModifier:         ll.opts.EnableAtModifier,
+		EnableNegativeOffset:     ll.opts.EnableNegativeOffset,
 	}
 
 	ll.queryEngine = NewEngine(opts)

promql/test_test.go

@@ -109,7 +109,7 @@ func TestLazyLoader_WithSamplesTill(t *testing.T) {
 	}
 
 	for _, c := range cases {
-		suite, err := NewLazyLoader(t, c.loadString)
+		suite, err := NewLazyLoader(t, c.loadString, LazyLoaderOpts{})
 		require.NoError(t, err)
 		defer suite.Close()