From 9d275c23de3c1d2c1e8c40bfb2f21eb86394d026 Mon Sep 17 00:00:00 2001
From: Julien <roidelapluie@o11y.eu>
Date: Fri, 4 Oct 2024 11:17:59 +0200
Subject: [PATCH] cmd/prometheus: Fix flakiness of QueryLogTest

Now we check that a rule execution has taken place.

This also reduces the time to run the rules tests from 45s to 25s.

Signed-off-by: Julien <roidelapluie@o11y.eu>
---
 cmd/prometheus/query_log_test.go | 55 ++++++++++++++++++++++++++++++--
 1 file changed, 52 insertions(+), 3 deletions(-)

diff --git a/cmd/prometheus/query_log_test.go b/cmd/prometheus/query_log_test.go
index 62e317bf8..3b00230cd 100644
--- a/cmd/prometheus/query_log_test.go
+++ b/cmd/prometheus/query_log_test.go
@@ -125,12 +125,61 @@ func (p *queryLogTest) query(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, 200, r.StatusCode)
 	case ruleOrigin:
-		time.Sleep(2 * time.Second)
+		// Poll the /api/v1/rules endpoint until a new rule evaluation is detected.
+		var lastEvalTime time.Time
+		for {
+			r, err := http.Get(fmt.Sprintf("http://%s:%d/api/v1/rules", p.host, p.port))
+			require.NoError(t, err)
+
+			rulesBody, err := io.ReadAll(r.Body)
+			require.NoError(t, err)
+			defer r.Body.Close()
+
+			// Parse the rules response to find the last evaluation time.
+			newEvalTime := parseLastEvaluation(rulesBody)
+			if newEvalTime.After(lastEvalTime) {
+				if !lastEvalTime.IsZero() {
+					break
+				}
+				lastEvalTime = newEvalTime
+			}
+
+			time.Sleep(100 * time.Millisecond)
+		}
 	default:
 		panic("can't query this origin")
 	}
 }
 
+// parseLastEvaluation extracts the last evaluation timestamp from the /api/v1/rules response.
+func parseLastEvaluation(rulesBody []byte) time.Time {
+	var ruleResponse struct {
+		Status string `json:"status"`
+		Data   struct {
+			Groups []struct {
+				Rules []struct {
+					LastEvaluation string `json:"lastEvaluation"`
+				} `json:"rules"`
+			} `json:"groups"`
+		} `json:"data"`
+	}
+
+	err := json.Unmarshal(rulesBody, &ruleResponse)
+	if err != nil {
+		return time.Time{}
+	}
+
+	for _, group := range ruleResponse.Data.Groups {
+		for _, rule := range group.Rules {
+			if evalTime, err := time.Parse(time.RFC3339Nano, rule.LastEvaluation); err == nil {
+				return evalTime
+			}
+		}
+	}
+
+	return time.Time{}
+}
+
 // queryString returns the expected queryString of a this test.
 func (p *queryLogTest) queryString() string {
 	switch p.origin {
@@ -322,7 +371,7 @@ func (p *queryLogTest) run(t *testing.T) {
 	if p.exactQueryCount() {
 		require.Len(t, ql, qc)
 	} else {
-		require.Greater(t, len(ql), qc, "no queries logged")
+		require.GreaterOrEqual(t, len(ql), qc, "no queries logged")
 	}
 	p.validateLastQuery(t, ql)
 	qc = len(ql)
@@ -353,7 +402,7 @@ func (p *queryLogTest) run(t *testing.T) {
 	if p.exactQueryCount() {
 		require.Len(t, ql, qc)
 	} else {
-		require.Greater(t, len(ql), qc, "no queries logged")
+		require.GreaterOrEqual(t, len(ql), qc, "no queries logged")
 	}
 	p.validateLastQuery(t, ql)