Rewrite TestTimeActiveStage tests (prometheus#3795)

This commit rewrites the existing TestTimeActiveStage unit tests
to have complete isolation between test cases. Before this change,
each test case affected the state of its subsequent tests.

The motivation behind this change is to make it easier to assert
that alerts have been marked as muted.

Signed-off-by: George Robinson <[email protected]>

Author: grobinson-grafana

notify/notify_test.go

@@ -19,6 +19,7 @@ import (
 	"fmt"
 	"io"
 	"reflect"
+	"strings"
 	"testing"
 	"time"
 
@@ -936,97 +937,96 @@ func TestTimeMuteStage(t *testing.T) {
 }
 
 func TestTimeActiveStage(t *testing.T) {
-	// Route mutes alerts inside business hours if it is an active time interval
-	muteIn := `
----
-- weekdays: ['monday:friday']
-  times:
-   - start_time: '00:00'
-     end_time: '09:00'
-   - start_time: '17:00'
-     end_time: '24:00'
-- weekdays: ['saturday', 'sunday']`
-
-	cases := []struct {
-		fireTime   string
-		labels     model.LabelSet
-		shouldMute bool
-	}{
-		{
-			// Friday during business hours
-			fireTime:   "01 Jan 21 09:00 +0000",
-			labels:     model.LabelSet{"mute": "me"},
-			shouldMute: true,
-		},
-		{
-			fireTime:   "02 Dec 20 16:59 +0000",
-			labels:     model.LabelSet{"mute": "me"},
-			shouldMute: true,
-		},
-		{
-			// Tuesday before 5pm
-			fireTime:   "01 Dec 20 16:59 +0000",
-			labels:     model.LabelSet{"mute": "me"},
-			shouldMute: true,
-		},
-		{
-			// Saturday
-			fireTime:   "17 Oct 20 10:00 +0000",
-			labels:     model.LabelSet{"foo": "bar"},
-			shouldMute: false,
-		},
-		{
-			// Wednesday before 9am
-			fireTime:   "14 Oct 20 05:00 +0000",
-			labels:     model.LabelSet{"dont": "mute"},
-			shouldMute: false,
-		},
-	}
-	var intervals []timeinterval.TimeInterval
-	err := yaml.Unmarshal([]byte(muteIn), &intervals)
+	sydney, err := time.LoadLocation("Australia/Sydney")
 	if err != nil {
-		t.Fatalf("Couldn't unmarshal time interval %s", err)
-	}
-	m := map[string][]timeinterval.TimeInterval{"test": intervals}
-	intervener := timeinterval.NewIntervener(m)
-	metrics := NewMetrics(prometheus.NewRegistry(), featurecontrol.NoopFlags{})
-	stage := NewTimeActiveStage(intervener, metrics)
-
-	outAlerts := []*types.Alert{}
-	nonMuteCount := 0
-	for _, tc := range cases {
-		now, err := time.Parse(time.RFC822Z, tc.fireTime)
-		if err != nil {
-			t.Fatalf("Couldn't parse fire time %s %s", tc.fireTime, err)
-		}
-		// Count alerts with shouldMute == false and compare to ensure none are muted incorrectly
-		if !tc.shouldMute {
-			nonMuteCount++
-		}
-		a := model.Alert{Labels: tc.labels}
-		alerts := []*types.Alert{{Alert: a}}
-		ctx := context.Background()
-		ctx = WithNow(ctx, now)
-		ctx = WithActiveTimeIntervals(ctx, []string{"test"})
-		ctx = WithMuteTimeIntervals(ctx, []string{})
+		t.Fatalf("Failed to load location Australia/Sydney: %s", err)
+	}
+	weekdays := map[string][]timeinterval.TimeInterval{
+		"weekdays": {{
+			Weekdays: []timeinterval.WeekdayRange{{
+				InclusiveRange: timeinterval.InclusiveRange{
+					Begin: 1, // Monday
+					End:   5, // Friday
+				},
+			}},
+			Times: []timeinterval.TimeRange{{
+				StartMinute: 540,  // 09:00
+				EndMinute:   1020, // 17:00
+			}},
+			Location: &timeinterval.Location{Location: sydney},
+		}},
+	}
+
+	tests := []struct {
+		name      string
+		intervals map[string][]timeinterval.TimeInterval
+		now       time.Time
+		alerts    []*types.Alert
+		mutedBy   []string
+	}{{
+		name:      "Should be muted outside working hours",
+		intervals: weekdays,
+		now:       time.Date(2024, 1, 1, 0, 0, 0, 0, sydney),
+		alerts:    []*types.Alert{{Alert: model.Alert{Labels: model.LabelSet{"foo": "bar"}}}},
+		mutedBy:   []string{"weekdays"},
+	}, {
+		name:      "Should not be muted during workings hours",
+		intervals: weekdays,
+		now:       time.Date(2024, 1, 1, 9, 0, 0, 0, sydney),
+		alerts:    []*types.Alert{{Alert: model.Alert{Labels: model.LabelSet{"foo": "bar"}}}},
+		mutedBy:   nil,
+	}, {
+		name:      "Should be muted during weekends",
+		intervals: weekdays,
+		now:       time.Date(2024, 1, 6, 10, 0, 0, 0, sydney),
+		alerts:    []*types.Alert{{Alert: model.Alert{Labels: model.LabelSet{"foo": "bar"}}}},
+		mutedBy:   []string{"weekdays"},
+	}, {
+		name:      "Should be muted at 12pm UTC",
+		intervals: weekdays,
+		now:       time.Date(2024, 1, 6, 10, 0, 0, 0, time.UTC),
+		alerts:    []*types.Alert{{Alert: model.Alert{Labels: model.LabelSet{"foo": "bar"}}}},
+		mutedBy:   []string{"weekdays"},
+	}}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			r := prometheus.NewRegistry()
+			metrics := NewMetrics(r, featurecontrol.NoopFlags{})
+			intervener := timeinterval.NewIntervener(test.intervals)
+			st := NewTimeActiveStage(intervener, metrics)
+
+			// Get the names of all time intervals for the context.
+			activeTimeIntervalNames := make([]string, 0, len(test.intervals))
+			for name := range test.intervals {
+				activeTimeIntervalNames = append(activeTimeIntervalNames, name)
+			}
 
-		_, out, err := stage.Exec(ctx, log.NewNopLogger(), alerts...)
-		if err != nil {
-			t.Fatalf("Unexpected error in time mute stage %s", err)
-		}
-		outAlerts = append(outAlerts, out...)
-	}
-	for _, alert := range outAlerts {
-		if _, ok := alert.Alert.Labels["mute"]; ok {
-			t.Fatalf("Expected alert to be muted %+v", alert.Alert)
-		}
-	}
-	if len(outAlerts) != nonMuteCount {
-		t.Fatalf("Expected %d alerts after time mute stage but got %d", nonMuteCount, len(outAlerts))
-	}
-	suppressed := int(prom_testutil.ToFloat64(metrics.numNotificationSuppressedTotal))
-	if (len(cases) - nonMuteCount) != suppressed {
-		t.Fatalf("Expected %d alerts counted in suppressed metric but got %d", (len(cases) - nonMuteCount), suppressed)
+			ctx := context.Background()
+			ctx = WithNow(ctx, test.now)
+			ctx = WithActiveTimeIntervals(ctx, activeTimeIntervalNames)
+			ctx = WithMuteTimeIntervals(ctx, nil)
+
+			_, active, err := st.Exec(ctx, log.NewNopLogger(), test.alerts...)
+			require.NoError(t, err)
+
+			if len(test.mutedBy) == 0 {
+				// All alerts should be active.
+				require.Equal(t, len(test.alerts), len(active))
+				// The metric for total suppressed notifications should not
+				// have been incremented, which means it will not be collected.
+				require.NoError(t, prom_testutil.GatherAndCompare(r, strings.NewReader("")))
+			} else {
+				// All alerts should be muted.
+				require.Empty(t, active)
+				// Gets the metric for total suppressed notifications.
+				require.NoError(t, prom_testutil.GatherAndCompare(r, strings.NewReader(fmt.Sprintf(`
+# HELP alertmanager_notifications_suppressed_total The total number of notifications suppressed for being silenced, inhibited, outside of active time intervals or within muted time intervals.
+# TYPE alertmanager_notifications_suppressed_total counter
+alertmanager_notifications_suppressed_total{reason="active_time_interval"} %d
+`, len(test.alerts)))))
+			}
+		})
 	}
 }