github.com/Azure/aad-pod-identity@v1.8.17/pkg/retry/retry_test.go (about)

     1  package retry
     2  
     3  import (
     4  	"errors"
     5  	"testing"
     6  
     7  	"github.com/stretchr/testify/assert"
     8  )
     9  
    10  func TestDo(t *testing.T) {
    11  	r := NewRetryClient(2, 0)
    12  	r.RegisterRetriableErrors("err1")
    13  
    14  	ran := 0
    15  	err := r.Do(func() error {
    16  		ran++
    17  		return nil
    18  	}, func(err error) bool {
    19  		return true
    20  	})
    21  	// Targeted function ran once since there is no error occurred
    22  	assert.Equal(t, 1, ran)
    23  	assert.NoError(t, err)
    24  
    25  	ran = 0
    26  	err = r.Do(func() error {
    27  		ran++
    28  		return errors.New("err1 occurred")
    29  	}, func(err error) bool {
    30  		return true
    31  	})
    32  	// Targeted function ran 3 times (1 initial run and 2 retries)
    33  	assert.Equal(t, 3, ran)
    34  	assert.Error(t, err)
    35  
    36  	ran = 0
    37  	err = r.Do(func() error {
    38  		ran++
    39  		return errors.New("err1 occurred")
    40  	}, func(err error) bool {
    41  		return false
    42  	})
    43  	// Targeted function ran once since shouldRetryFunc returned false
    44  	assert.Equal(t, 1, ran)
    45  	assert.Error(t, err)
    46  
    47  	ran = 0
    48  	err = r.Do(func() error {
    49  		ran++
    50  		return errors.New("err2 occurred")
    51  	}, func(err error) bool {
    52  		return true
    53  	})
    54  	// Targeted function only ran once since err2 was not registered
    55  	assert.Equal(t, 1, ran)
    56  	assert.Error(t, err)
    57  }