Apply changes to unit tests to other optimization passes

Author: zenador

pkg/streamingpromql/optimize/ast/collapse_constants_test.go

@@ -1,11 +1,14 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 
-package ast
+package ast_test
 
 import (
 	"context"
 	"testing"
 
+	"github.com/grafana/mimir/pkg/streamingpromql/optimize"
+	"github.com/grafana/mimir/pkg/streamingpromql/optimize/ast"
+	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/prometheus/promql/parser"
 	"github.com/stretchr/testify/require"
 )
@@ -51,15 +54,15 @@ func TestCollapseConstants(t *testing.T) {
 		`"abc"`:         `"abc"`,
 	}
 
-	collapseConstants := &CollapseConstants{}
+	ctx := context.Background()
+	collapseConstants := &ast.CollapseConstants{}
 
 	for input, expected := range testCases {
 		t.Run(input, func(t *testing.T) {
-			expr, err := parser.ParseExpr(input)
-			require.NoError(t, err)
+			_, _, result := getOutputFromASTOptimizationPassWithQueryPlan(t, ctx, input, func(prometheus.Registerer) optimize.ASTOptimizationPass {
+				return collapseConstants
+			})
 
-			result, err := collapseConstants.Apply(context.Background(), expr)
-			require.NoError(t, err)
 			require.Equal(t, expected, result.String())
 
 			// Check for unnecessary unary expressions.

pkg/streamingpromql/optimize/ast/sort_labels_and_matchers_test.go

@@ -1,13 +1,16 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 
-package ast
+package ast_test
 
 import (
 	"context"
 	"fmt"
 	"testing"
 	"time"
 
+	"github.com/grafana/mimir/pkg/streamingpromql/optimize"
+	"github.com/grafana/mimir/pkg/streamingpromql/optimize/ast"
+	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/promql/parser"
 	"github.com/prometheus/prometheus/promql/parser/posrange"
@@ -56,15 +59,15 @@ func TestSortLabelsAndMatchers_AggregateAndBinaryExpressions(t *testing.T) {
 		"2":          "2",
 	}
 
-	sortLabelsAndMatchers := &SortLabelsAndMatchers{}
+	ctx := context.Background()
+	sortLabelsAndMatchers := &ast.SortLabelsAndMatchers{}
 
 	for input, expected := range testCases {
 		t.Run(input, func(t *testing.T) {
-			expr, err := parser.ParseExpr(input)
-			require.NoError(t, err)
+			_, _, result := getOutputFromASTOptimizationPassWithQueryPlan(t, ctx, input, func(prometheus.Registerer) optimize.ASTOptimizationPass {
+				return sortLabelsAndMatchers
+			})
 
-			result, err := sortLabelsAndMatchers.Apply(context.Background(), expr)
-			require.NoError(t, err)
 			require.Equal(t, expected, result.String())
 		})
 	}
@@ -154,14 +157,13 @@ func TestSortLabelsAndMatchers_Selectors(t *testing.T) {
 		},
 	}
 
-	sortLabelsAndMatchers := &SortLabelsAndMatchers{}
+	ctx := context.Background()
+	sortLabelsAndMatchers := &ast.SortLabelsAndMatchers{}
 
 	run := func(t *testing.T, input string, expected parser.Expr) {
-		expr, err := parser.ParseExpr(input)
-		require.NoError(t, err)
-
-		result, err := sortLabelsAndMatchers.Apply(context.Background(), expr)
-		require.NoError(t, err)
+		_, _, result := getOutputFromASTOptimizationPassWithQueryPlan(t, ctx, input, func(prometheus.Registerer) optimize.ASTOptimizationPass {
+			return sortLabelsAndMatchers
+		})
 
 		// Compare the expected and result expressions formatted as strings for clearer diffs.
 		// Note that we can't use the VectorSelector and MatrixSelector String() methods because these