github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/module/retrymiddleware/after_consecutive_failures_test.go (about)

     1  package retrymiddleware
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"fmt"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/sethvargo/go-retry"
    11  	"github.com/stretchr/testify/require"
    12  
    13  	"github.com/onflow/flow-go/model/messages"
    14  	"github.com/onflow/flow-go/module/mock"
    15  )
    16  
    17  // TestAfterConsecutiveFailures test that the middleware executes the onConsecutiveFailures func as expected
    18  func TestAfterConsecutiveFailures(t *testing.T) {
    19  	a := &mock.DKGContractClient{}
    20  	b := &mock.DKGContractClient{}
    21  	c := &mock.DKGContractClient{}
    22  
    23  	msg := messages.BroadcastDKGMessage{}
    24  
    25  	a.On("Broadcast", msg).
    26  		Return(errors.New("error from a")).
    27  		Twice()
    28  
    29  	b.On("Broadcast", msg).
    30  		Return(errors.New("error from b")).
    31  		Twice()
    32  
    33  	c.On("Broadcast", msg).
    34  		Return(errors.New("error from c")).
    35  		Twice()
    36  
    37  	clients := []*mock.DKGContractClient{a, b, c}
    38  
    39  	// every 2 failures we will update our dkgContractClient
    40  	maxConsecutiveRetries := 2
    41  
    42  	backoff := retry.NewConstant(1000 * time.Millisecond)
    43  	backoff = retry.WithMaxRetries(5, backoff)
    44  
    45  	// after 2 consecutive failures fallback to next DKGContractClient
    46  	clientIndex := 0
    47  	dkgContractClient := clients[clientIndex]
    48  	afterConsecutiveFailures := AfterConsecutiveFailures(maxConsecutiveRetries, backoff, func(totalAttempts int) {
    49  		if clientIndex == len(clients)-1 {
    50  			clientIndex = 0
    51  		} else {
    52  			clientIndex++
    53  		}
    54  		dkgContractClient = clients[clientIndex]
    55  	})
    56  
    57  	err := retry.Do(context.Background(), afterConsecutiveFailures, func(ctx context.Context) error {
    58  		err := dkgContractClient.Broadcast(msg)
    59  		if err != nil {
    60  			fmt.Printf("error: %s\n", err)
    61  		}
    62  		return retry.RetryableError(err)
    63  	})
    64  	require.Error(t, err)
    65  
    66  	a.AssertExpectations(t)
    67  	b.AssertExpectations(t)
    68  	c.AssertExpectations(t)
    69  }