Fixing Flaky AM test (#4865)

Signed-off-by: Alan Protasio <[email protected]>

Signed-off-by: Alan Protasio <[email protected]>

Author: alanprot

pkg/alertmanager/distributor_test.go

@@ -19,6 +19,7 @@ import (
 	"github.com/stretchr/testify/require"
 	"github.com/weaveworks/common/httpgrpc"
 	"github.com/weaveworks/common/user"
+	"go.uber.org/atomic"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/health/grpc_health_v1"
 
@@ -333,11 +334,9 @@ func TestDistributor_IsPathSupported(t *testing.T) {
 
 func prepare(t *testing.T, numAM, numHappyAM, replicationFactor int, responseBody []byte) (*Distributor, []*mockAlertmanager, func()) {
 	ams := []*mockAlertmanager{}
-	for i := 0; i < numHappyAM; i++ {
-		ams = append(ams, newMockAlertmanager(i, true, responseBody))
-	}
-	for i := numHappyAM; i < numAM; i++ {
-		ams = append(ams, newMockAlertmanager(i, false, responseBody))
+	remainingFailure := atomic.NewInt32(int32(numAM - numHappyAM))
+	for i := 0; i < numAM; i++ {
+		ams = append(ams, newMockAlertmanager(i, remainingFailure, responseBody))
 	}
 
 	// Use a real ring with a mock KV store to test ring RF logic.
@@ -399,15 +398,15 @@ type mockAlertmanager struct {
 	receivedRequests map[string]map[int]int
 	mtx              sync.Mutex
 	myAddr           string
-	happy            bool
+	remainingError   *atomic.Int32
 	responseBody     []byte
 }
 
-func newMockAlertmanager(idx int, happy bool, responseBody []byte) *mockAlertmanager {
+func newMockAlertmanager(idx int, remainingError *atomic.Int32, responseBody []byte) *mockAlertmanager {
 	return &mockAlertmanager{
 		receivedRequests: make(map[string]map[int]int),
 		myAddr:           fmt.Sprintf("127.0.0.1:%05d", 10000+idx),
-		happy:            happy,
+		remainingError:   remainingError,
 		responseBody:     responseBody,
 	}
 }
@@ -427,7 +426,7 @@ func (am *mockAlertmanager) HandleRequest(_ context.Context, in *httpgrpc.HTTPRe
 		am.receivedRequests[path] = m
 	}
 
-	if am.happy {
+	if am.remainingError.Dec() < 0 {
 		m[http.StatusOK]++
 		return &httpgrpc.HTTPResponse{
 			Code: http.StatusOK,