git.sr.ht/~pingoo/stdx@v0.0.0-20240218134121-094174641f6e/retry/examples/delay_based_on_error_test.go (about)

     1  // This test delay is based on kind of error
     2  // e.g. HTTP response [Retry-After](https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Retry-After)
     3  package retry_test
     4  
     5  import (
     6  	"fmt"
     7  	"io/ioutil"
     8  	"net/http"
     9  	"net/http/httptest"
    10  	"testing"
    11  	"time"
    12  
    13  	"git.sr.ht/~pingoo/stdx/retry"
    14  	"github.com/stretchr/testify/assert"
    15  )
    16  
    17  type RetryAfterError struct {
    18  	response http.Response
    19  }
    20  
    21  func (err RetryAfterError) Error() string {
    22  	return fmt.Sprintf(
    23  		"Request to %s fail %s (%d)",
    24  		err.response.Request.RequestURI,
    25  		err.response.Status,
    26  		err.response.StatusCode,
    27  	)
    28  }
    29  
    30  type SomeOtherError struct {
    31  	err        string
    32  	retryAfter time.Duration
    33  }
    34  
    35  func (err SomeOtherError) Error() string {
    36  	return err.err
    37  }
    38  
    39  func TestCustomRetryFunctionBasedOnKindOfError(t *testing.T) {
    40  	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    41  		fmt.Fprintln(w, "hello")
    42  	}))
    43  	defer ts.Close()
    44  
    45  	var body []byte
    46  
    47  	err := retry.Do(
    48  		func() error {
    49  			resp, err := http.Get(ts.URL)
    50  
    51  			if err == nil {
    52  				defer func() {
    53  					if err := resp.Body.Close(); err != nil {
    54  						panic(err)
    55  					}
    56  				}()
    57  				body, err = ioutil.ReadAll(resp.Body)
    58  			}
    59  
    60  			return err
    61  		},
    62  		retry.DelayType(func(n uint, err error, config *retry.Config) time.Duration {
    63  			switch e := err.(type) {
    64  			case RetryAfterError:
    65  				if t, err := parseRetryAfter(e.response.Header.Get("Retry-After")); err == nil {
    66  					return time.Until(t)
    67  				}
    68  			case SomeOtherError:
    69  				return e.retryAfter
    70  			}
    71  
    72  			//default is backoffdelay
    73  			return retry.BackOffDelay(n, err, config)
    74  		}),
    75  	)
    76  
    77  	assert.NoError(t, err)
    78  	assert.NotEmpty(t, body)
    79  }
    80  
    81  // use https://github.com/aereal/go-httpretryafter instead
    82  func parseRetryAfter(_ string) (time.Time, error) {
    83  	return time.Now().Add(1 * time.Second), nil
    84  }