github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/pkg/querier/queryrange/retry_test.go (about)

     1  package queryrange
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	fmt "fmt"
     7  	"net/http"
     8  	"testing"
     9  
    10  	"github.com/go-kit/log"
    11  	"github.com/stretchr/testify/require"
    12  	"github.com/weaveworks/common/httpgrpc"
    13  	"go.uber.org/atomic"
    14  )
    15  
    16  func TestRetry(t *testing.T) {
    17  	var try atomic.Int32
    18  
    19  	for _, tc := range []struct {
    20  		name    string
    21  		handler Handler
    22  		resp    Response
    23  		err     error
    24  	}{
    25  		{
    26  			name: "retry failures",
    27  			handler: HandlerFunc(func(_ context.Context, req Request) (Response, error) {
    28  				if try.Inc() == 5 {
    29  					return &PrometheusResponse{Status: "Hello World"}, nil
    30  				}
    31  				return nil, fmt.Errorf("fail")
    32  			}),
    33  			resp: &PrometheusResponse{Status: "Hello World"},
    34  		},
    35  		{
    36  			name: "don't retry 400s",
    37  			handler: HandlerFunc(func(_ context.Context, req Request) (Response, error) {
    38  				return nil, httpgrpc.Errorf(http.StatusBadRequest, "Bad Request")
    39  			}),
    40  			err: httpgrpc.Errorf(http.StatusBadRequest, "Bad Request"),
    41  		},
    42  		{
    43  			name: "retry 500s",
    44  			handler: HandlerFunc(func(_ context.Context, req Request) (Response, error) {
    45  				return nil, httpgrpc.Errorf(http.StatusInternalServerError, "Internal Server Error")
    46  			}),
    47  			err: httpgrpc.Errorf(http.StatusInternalServerError, "Internal Server Error"),
    48  		},
    49  		{
    50  			name: "last error",
    51  			handler: HandlerFunc(func(_ context.Context, req Request) (Response, error) {
    52  				if try.Inc() == 5 {
    53  					return nil, httpgrpc.Errorf(http.StatusBadRequest, "Bad Request")
    54  				}
    55  				return nil, httpgrpc.Errorf(http.StatusInternalServerError, "Internal Server Error")
    56  			}),
    57  			err: httpgrpc.Errorf(http.StatusBadRequest, "Bad Request"),
    58  		},
    59  	} {
    60  		t.Run(tc.name, func(t *testing.T) {
    61  			try.Store(0)
    62  			h := NewRetryMiddleware(log.NewNopLogger(), 5, nil).Wrap(tc.handler)
    63  			resp, err := h.Do(context.Background(), nil)
    64  			require.Equal(t, tc.err, err)
    65  			require.Equal(t, tc.resp, resp)
    66  		})
    67  	}
    68  }
    69  
    70  func Test_RetryMiddlewareCancel(t *testing.T) {
    71  	var try atomic.Int32
    72  	ctx, cancel := context.WithCancel(context.Background())
    73  	cancel()
    74  	_, err := NewRetryMiddleware(log.NewNopLogger(), 5, nil).Wrap(
    75  		HandlerFunc(func(c context.Context, r Request) (Response, error) {
    76  			try.Inc()
    77  			return nil, ctx.Err()
    78  		}),
    79  	).Do(ctx, nil)
    80  	require.Equal(t, int32(0), try.Load())
    81  	require.Equal(t, ctx.Err(), err)
    82  
    83  	ctx, cancel = context.WithCancel(context.Background())
    84  	_, err = NewRetryMiddleware(log.NewNopLogger(), 5, nil).Wrap(
    85  		HandlerFunc(func(c context.Context, r Request) (Response, error) {
    86  			try.Inc()
    87  			cancel()
    88  			return nil, errors.New("failed")
    89  		}),
    90  	).Do(ctx, nil)
    91  	require.Equal(t, int32(1), try.Load())
    92  	require.Equal(t, ctx.Err(), err)
    93  }