github.com/thanos-io/thanos@v0.32.5/internal/cortex/querier/queryrange/retry_test.go (about)

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