provider/mem: add test detecting dropped alerts

Signed-off-by: Simon Pasquier <spasquie@redhat.com>

provider/mem/mem_test.go

@@ -161,91 +161,82 @@ func TestAlertsPut(t *testing.T) {
 
 func TestAlertsSubscribe(t *testing.T) {
 	marker := types.NewMarker(prometheus.NewRegistry())
-	alerts, err := NewAlerts(context.Background(), marker, 30*time.Minute, log.NewNopLogger())
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	alerts, err := NewAlerts(ctx, marker, 30*time.Minute, log.NewNopLogger())
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	// add alert1 to validate if pending alerts will be sent
+	// Add alert1 to validate if pending alerts will be sent.
 	if err := alerts.Put(alert1); err != nil {
 		t.Fatalf("Insert failed: %s", err)
 	}
 
-	var wg sync.WaitGroup
-	wg.Add(2)
-	fatalc := make(chan string, 2)
+	expectedAlerts := map[model.Fingerprint]*types.Alert{
+		alert1.Fingerprint(): alert1,
+		alert2.Fingerprint(): alert2,
+		alert3.Fingerprint(): alert3,
+	}
 
-	iterator1 := alerts.Subscribe()
-	iterator2 := alerts.Subscribe()
+	// Start many consumers and make sure that each receives all the subsequent alerts.
+	var (
+		nb     = 100
+		fatalc = make(chan string, nb)
+		wg     sync.WaitGroup
+	)
+	wg.Add(nb)
+	for i := 0; i < nb; i++ {
+		go func(i int) {
+			defer wg.Done()
 
-	go func() {
-		defer wg.Done()
-		expectedAlerts := map[model.Fingerprint]*types.Alert{
-			alert1.Fingerprint(): alert1,
-			alert2.Fingerprint(): alert2,
-			alert3.Fingerprint(): alert3,
-		}
+			it := alerts.Subscribe()
+			defer it.Close()
 
-		for i := 0; i < 3; i++ {
-			actual := <-iterator1.Next()
-			expected := expectedAlerts[actual.Fingerprint()]
-			if !alertsEqual(actual, expected) {
-				fatalc <- fmt.Sprintf("Unexpected alert (iterator1)\n%s", pretty.Compare(actual, expected))
-				return
+			received := make(map[model.Fingerprint]struct{})
+			for {
+				select {
+				case got, ok := <-it.Next():
+					if !ok {
+						fatalc <- fmt.Sprintf("Iterator %d closed", i)
+						return
+					}
+					if it.Err() != nil {
+						fatalc <- fmt.Sprintf("Iterator %d: %v", i, it.Err())
+						return
+					}
+					expected := expectedAlerts[got.Fingerprint()]
+					if !alertsEqual(got, expected) {
+						fatalc <- fmt.Sprintf("Unexpected alert (iterator %d)\n%s", i, pretty.Compare(got, expected))
+						return
+					}
+					received[got.Fingerprint()] = struct{}{}
+					if len(received) == len(expectedAlerts) {
+						return
+					}
+				case <-time.After(5 * time.Second):
+					fatalc <- fmt.Sprintf("Unexpected number of alerts for iterator %d, got: %d, expected: %d", i, len(received), len(expectedAlerts))
+					return
+				}
 			}
+		}(i)
+	}
 
-			delete(expectedAlerts, actual.Fingerprint())
-		}
-
-		if len(expectedAlerts) != 0 {
-			fatalc <- fmt.Sprintf("Unexpected number of alerts (iterator1): %d", len(expectedAlerts))
-		}
-	}()
-
-	go func() {
-		defer wg.Done()
-		expectedAlerts := map[model.Fingerprint]*types.Alert{
-			alert1.Fingerprint(): alert1,
-			alert2.Fingerprint(): alert2,
-			alert3.Fingerprint(): alert3,
-		}
-
-		for i := 0; i < 3; i++ {
-			actual := <-iterator2.Next()
-			expected := expectedAlerts[actual.Fingerprint()]
-			if !alertsEqual(actual, expected) {
-				t.Errorf("Unexpected alert")
-				fatalc <- fmt.Sprintf("Unexpected alert (iterator2)\n%s", pretty.Compare(actual, expected))
-			}
-
-			delete(expectedAlerts, actual.Fingerprint())
-		}
-
-		if len(expectedAlerts) != 0 {
-			fatalc <- fmt.Sprintf("Unexpected number of alerts (iterator2): %d", len(expectedAlerts))
-		}
-	}()
-
-	go func() {
-		wg.Wait()
-		close(fatalc)
-	}()
-
+	// Add more alerts that should be received by the subscribers.
 	if err := alerts.Put(alert2); err != nil {
 		t.Fatalf("Insert failed: %s", err)
 	}
-
 	if err := alerts.Put(alert3); err != nil {
 		t.Fatalf("Insert failed: %s", err)
 	}
 
+	wg.Wait()
+	close(fatalc)
 	fatal, ok := <-fatalc
 	if ok {
 		t.Fatalf(fatal)
 	}
-
-	iterator1.Close()
-	iterator2.Close()
 }
 
 func TestAlertsGetPending(t *testing.T) {