prometheus#3513: Rewrite TestTimeMuteStage tests (prometheus#3794)

Author: grobinson-grafana

notify/notify_test.go

@@ -28,7 +28,6 @@ import (
 	prom_testutil "github.com/prometheus/client_golang/prometheus/testutil"
 	"github.com/prometheus/common/model"
 	"github.com/stretchr/testify/require"
-	"gopkg.in/yaml.v2"
 
 	"github.com/prometheus/alertmanager/featurecontrol"
 	"github.com/prometheus/alertmanager/nflog"
@@ -822,117 +821,107 @@ func TestMuteStageWithSilences(t *testing.T) {
 }
 
 func TestTimeMuteStage(t *testing.T) {
-	// Route mutes alerts outside business hours in November, using the +1100 timezone.
-	muteIn := `
----
-- weekdays: ['monday:friday']
-  location: 'Australia/Sydney'
-  months: ['November']
-  times:
-   - start_time: '00:00'
-     end_time: '09:00'
-   - start_time: '17:00'
-     end_time: '24:00'
-- weekdays: ['saturday', 'sunday']
-  months: ['November']
-  location: 'Australia/Sydney'`
-
-	cases := []struct {
-		fireTime   string
-		labels     model.LabelSet
-		shouldMute bool
-	}{
-		{
-			// Friday during business hours
-			fireTime:   "19 Nov 21 13:00 +1100",
-			labels:     model.LabelSet{"foo": "bar"},
-			shouldMute: false,
-		},
-		{
-			// Tuesday before 5pm
-			fireTime:   "16 Nov 21 16:59 +1100",
-			labels:     model.LabelSet{"dont": "mute"},
-			shouldMute: false,
-		},
-		{
-			// Saturday
-			fireTime:   "20 Nov 21 10:00 +1100",
-			labels:     model.LabelSet{"mute": "me"},
-			shouldMute: true,
-		},
-		{
-			// Wednesday before 9am
-			fireTime:   "17 Nov 21 05:00 +1100",
-			labels:     model.LabelSet{"mute": "me"},
-			shouldMute: true,
-		},
-		{
-			// Ensure comparisons with other time zones work as expected.
-			fireTime:   "14 Nov 21 20:00 +0900",
-			labels:     model.LabelSet{"mute": "kst"},
-			shouldMute: true,
-		},
-		{
-			fireTime:   "14 Nov 21 21:30 +0000",
-			labels:     model.LabelSet{"mute": "utc"},
-			shouldMute: true,
-		},
-		{
-			fireTime:   "15 Nov 22 14:30 +0900",
-			labels:     model.LabelSet{"kst": "dont_mute"},
-			shouldMute: false,
-		},
-		{
-			fireTime:   "15 Nov 21 02:00 -0500",
-			labels:     model.LabelSet{"mute": "0500"},
-			shouldMute: true,
-		},
-	}
-	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 := NewTimeMuteStage(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{})
-		ctx = WithMuteTimeIntervals(ctx, []string{"test"})
-
-		_, 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)
-		}
+		t.Fatalf("Failed to load location Australia/Sydney: %s", err)
 	}
-	if len(outAlerts) != nonMuteCount {
-		t.Fatalf("Expected %d alerts after time mute stage but got %d", nonMuteCount, len(outAlerts))
+	eveningsAndWeekends := map[string][]timeinterval.TimeInterval{
+		"evenings": {{
+			Weekdays: []timeinterval.WeekdayRange{{
+				InclusiveRange: timeinterval.InclusiveRange{
+					Begin: 1, // Monday
+					End:   5, // Friday
+				},
+			}},
+			Times: []timeinterval.TimeRange{{
+				StartMinute: 0,   // 00:00
+				EndMinute:   540, // 09:00
+			}, {
+				StartMinute: 1020, // 17:00
+				EndMinute:   1440, // 24:00
+			}},
+			Location: &timeinterval.Location{Location: sydney},
+		}},
+		"weekends": {{
+			Weekdays: []timeinterval.WeekdayRange{{
+				InclusiveRange: timeinterval.InclusiveRange{Begin: 6, End: 6}, // Saturday
+			}, {
+				InclusiveRange: timeinterval.InclusiveRange{Begin: 0, End: 0}, // Sunday
+			}},
+			Location: &timeinterval.Location{Location: sydney},
+		}},
 	}
-	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)
+
+	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: eveningsAndWeekends,
+		now:       time.Date(2024, 1, 1, 0, 0, 0, 0, sydney),
+		alerts:    []*types.Alert{{Alert: model.Alert{Labels: model.LabelSet{"foo": "bar"}}}},
+		mutedBy:   []string{"evenings"},
+	}, {
+		name:      "Should not be muted during workings hours",
+		intervals: eveningsAndWeekends,
+		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: eveningsAndWeekends,
+		now:       time.Date(2024, 1, 6, 10, 0, 0, 0, sydney),
+		alerts:    []*types.Alert{{Alert: model.Alert{Labels: model.LabelSet{"foo": "bar"}}}},
+		mutedBy:   []string{"weekends"},
+	}, {
+		name:      "Should be muted at 12pm UTC",
+		intervals: eveningsAndWeekends,
+		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{"evenings"},
+	}}
+
+	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 := NewTimeMuteStage(intervener, metrics)
+
+			// Get the names of all time intervals for the context.
+			muteTimeIntervalNames := make([]string, 0, len(test.intervals))
+			for name := range test.intervals {
+				muteTimeIntervalNames = append(muteTimeIntervalNames, name)
+			}
+
+			ctx := context.Background()
+			ctx = WithNow(ctx, test.now)
+			ctx = WithActiveTimeIntervals(ctx, nil)
+			ctx = WithMuteTimeIntervals(ctx, muteTimeIntervalNames)
+
+			_, 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="mute_time_interval"} %d
+`, len(test.alerts)))))
+			}
+		})
 	}
 }