From 77c784ac934e7a0d7f78f198d38b0adeec12e1ee Mon Sep 17 00:00:00 2001
From: David Leadbeater <dgl@dgl.cx>
Date: Mon, 21 Sep 2020 17:15:34 +0100
Subject: [PATCH] Ensure alert rules are marked as restored in unit tests
 (#7661)

This makes sure the ALERTS timeseries is created when unit testing
alerting rules.

Signed-off-by: David Leadbeater <dgl@dgl.cx>
---
 cmd/promtool/testdata/unittest.yml |  8 ++++++++
 cmd/promtool/unittest.go           | 10 ++++++++++
 2 files changed, 18 insertions(+)

diff --git a/cmd/promtool/testdata/unittest.yml b/cmd/promtool/testdata/unittest.yml
index f8fb2d758..e25a09da5 100644
--- a/cmd/promtool/testdata/unittest.yml
+++ b/cmd/promtool/testdata/unittest.yml
@@ -8,6 +8,14 @@ tests:
     input_series:
       - series: 'up{job="prometheus", instance="localhost:9090"}'
         values: "0+0x1440"
+
+    promql_expr_test:
+      - expr: count(ALERTS) by (alertname, alertstate)
+        eval_time: 4m
+        exp_samples:
+        - labels: '{alertname="InstanceDown",alertstate="pending"}'
+          value: 1
+
     alert_rule_test:
       - eval_time: 1d
         alertname: InstanceDown
diff --git a/cmd/promtool/unittest.go b/cmd/promtool/unittest.go
index 42368b391..cbdb8931e 100644
--- a/cmd/promtool/unittest.go
+++ b/cmd/promtool/unittest.go
@@ -221,6 +221,16 @@ func (tg *testGroup) test(mint, maxt time.Time, evalInterval time.Duration, grou
 	// Current index in alertEvalTimes what we are looking at.
 	curr := 0
 
+	for _, g := range groups {
+		for _, r := range g.Rules() {
+			if alertRule, ok := r.(*rules.AlertingRule); ok {
+				// Mark alerting rules as restored, to ensure the ALERTS timeseries is
+				// created when they run.
+				alertRule.SetRestored(true)
+			}
+		}
+	}
+
 	var errs []error
 	for ts := mint; ts.Before(maxt); ts = ts.Add(evalInterval) {
 		// Collects the alerts asked for unit testing.