github.com/ydb-platform/ydb-go-sdk/v3@v3.57.0/retry/retry_test.go (about)

     1  package retry
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/stretchr/testify/require"
    10  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
    11  	grpcCodes "google.golang.org/grpc/codes"
    12  	grpcStatus "google.golang.org/grpc/status"
    13  
    14  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
    15  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
    16  )
    17  
    18  func TestRetryModes(t *testing.T) {
    19  	for _, idempotentType := range []idempotency{
    20  		idempotent,
    21  		nonIdempotent,
    22  	} {
    23  		t.Run(idempotentType.String(), func(t *testing.T) {
    24  			for _, tt := range errsToCheck {
    25  				t.Run(tt.err.Error(), func(t *testing.T) {
    26  					m := Check(tt.err)
    27  					if m.MustRetry(true) != tt.canRetry[idempotent] {
    28  						t.Errorf(
    29  							"unexpected must retry idempotent operation status: %v, want: %v",
    30  							m.MustRetry(true),
    31  							tt.canRetry[idempotent],
    32  						)
    33  					}
    34  					if m.MustRetry(false) != tt.canRetry[nonIdempotent] {
    35  						t.Errorf(
    36  							"unexpected must retry non-idempotent operation status: %v, want: %v",
    37  							m.MustRetry(false),
    38  							tt.canRetry[nonIdempotent],
    39  						)
    40  					}
    41  					if m.BackoffType() != tt.backoff {
    42  						t.Errorf(
    43  							"unexpected backoff status: %v, want: %v",
    44  							m.BackoffType(),
    45  							tt.backoff,
    46  						)
    47  					}
    48  					if m.MustDeleteSession() != tt.deleteSession {
    49  						t.Errorf(
    50  							"unexpected delete session status: %v, want: %v",
    51  							m.MustDeleteSession(),
    52  							tt.deleteSession,
    53  						)
    54  					}
    55  				})
    56  			}
    57  		})
    58  	}
    59  }
    60  
    61  type CustomError struct {
    62  	Err error
    63  }
    64  
    65  func (e *CustomError) Error() string {
    66  	return fmt.Sprintf("custom error: %v", e.Err)
    67  }
    68  
    69  func (e *CustomError) Unwrap() error {
    70  	return e.Err
    71  }
    72  
    73  func TestRetryWithCustomErrors(t *testing.T) {
    74  	var (
    75  		limit = 10
    76  		ctx   = context.Background()
    77  	)
    78  	for _, tt := range []struct {
    79  		error     error
    80  		retriable bool
    81  	}{
    82  		{
    83  			error: &CustomError{
    84  				Err: RetryableError(
    85  					fmt.Errorf("custom error"),
    86  					WithDeleteSession(),
    87  				),
    88  			},
    89  			retriable: true,
    90  		},
    91  		{
    92  			error: &CustomError{
    93  				Err: xerrors.Operation(
    94  					xerrors.WithStatusCode(Ydb.StatusIds_BAD_SESSION),
    95  				),
    96  			},
    97  			retriable: true,
    98  		},
    99  		{
   100  			error: &CustomError{
   101  				Err: fmt.Errorf(
   102  					"wrapped error: %w",
   103  					xerrors.Operation(
   104  						xerrors.WithStatusCode(Ydb.StatusIds_BAD_SESSION),
   105  					),
   106  				),
   107  			},
   108  			retriable: true,
   109  		},
   110  		{
   111  			error: &CustomError{
   112  				Err: fmt.Errorf(
   113  					"wrapped error: %w",
   114  					xerrors.Operation(
   115  						xerrors.WithStatusCode(Ydb.StatusIds_UNAUTHORIZED),
   116  					),
   117  				),
   118  			},
   119  			retriable: false,
   120  		},
   121  	} {
   122  		t.Run(tt.error.Error(), func(t *testing.T) {
   123  			i := 0
   124  			err := Retry(ctx, func(ctx context.Context) error {
   125  				i++
   126  				if i < limit {
   127  					return tt.error
   128  				}
   129  
   130  				return nil
   131  			})
   132  			if tt.retriable {
   133  				if i != limit {
   134  					t.Fatalf("unexpected i: %d, queryErr: %v", i, err)
   135  				}
   136  			} else {
   137  				if i != 1 {
   138  					t.Fatalf("unexpected i: %d, queryErr: %v", i, err)
   139  				}
   140  			}
   141  		})
   142  	}
   143  }
   144  
   145  func TestRetryTransportDeadlineExceeded(t *testing.T) {
   146  	cancelCounterValue := 5
   147  	for _, code := range []grpcCodes.Code{
   148  		grpcCodes.DeadlineExceeded,
   149  		grpcCodes.Canceled,
   150  	} {
   151  		counter := 0
   152  		ctx, cancel := xcontext.WithTimeout(context.Background(), time.Hour)
   153  		err := Retry(ctx, func(ctx context.Context) error {
   154  			counter++
   155  			if !(counter < cancelCounterValue) {
   156  				cancel()
   157  			}
   158  
   159  			return xerrors.Transport(grpcStatus.Error(code, ""))
   160  		}, WithIdempotent(true))
   161  		require.ErrorIs(t, err, context.Canceled)
   162  		require.Equal(t, cancelCounterValue, counter)
   163  	}
   164  }
   165  
   166  func TestRetryTransportCancelled(t *testing.T) {
   167  	cancelCounterValue := 5
   168  	for _, code := range []grpcCodes.Code{
   169  		grpcCodes.DeadlineExceeded,
   170  		grpcCodes.Canceled,
   171  	} {
   172  		counter := 0
   173  		ctx, cancel := xcontext.WithCancel(context.Background())
   174  		err := Retry(ctx, func(ctx context.Context) error {
   175  			counter++
   176  			if !(counter < cancelCounterValue) {
   177  				cancel()
   178  			}
   179  
   180  			return xerrors.Transport(grpcStatus.Error(code, ""))
   181  		}, WithIdempotent(true))
   182  		require.ErrorIs(t, err, context.Canceled)
   183  		require.Equal(t, cancelCounterValue, counter)
   184  	}
   185  }