From a6e427660aa7e6b360049f8bf6990c3bdb80a9aa Mon Sep 17 00:00:00 2001
From: Bryan Boreham <bjboreham@gmail.com>
Date: Sat, 4 May 2024 14:22:28 +0100
Subject: [PATCH] test: check for @-modifier without using engine internals

Signed-off-by: Bryan Boreham <bjboreham@gmail.com>
---
 promql/test.go | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/promql/test.go b/promql/test.go
index b982c33c3..2cfcfce6f 100644
--- a/promql/test.go
+++ b/promql/test.go
@@ -676,8 +676,7 @@ func atModifierTestCases(exprStr string, evalTime time.Time) ([]atModifierTestCa
 	// If there is a subquery, then the selectors inside it don't get the @ timestamp.
 	// If any selector already has the @ timestamp set, then it is untouched.
 	parser.Inspect(expr, func(node parser.Node, path []parser.Node) error {
-		_, _, subqTs := subqueryTimes(path)
-		if subqTs != nil {
+		if hasAtModifier(path) {
 			// There is a subquery with timestamp in the path,
 			// hence don't change any timestamps further.
 			return nil
@@ -727,6 +726,17 @@ func atModifierTestCases(exprStr string, evalTime time.Time) ([]atModifierTestCa
 	return testCases, nil
 }
 
+func hasAtModifier(path []parser.Node) bool {
+	for _, node := range path {
+		if n, ok := node.(*parser.SubqueryExpr); ok {
+			if n.Timestamp != nil {
+				return true
+			}
+		}
+	}
+	return false
+}
+
 // exec processes a single step of the test.
 func (t *test) exec(tc testCommand, engine QueryEngine) error {
 	switch cmd := tc.(type) {