retry_test.go

Documentation: cloud.google.com/go/bigquery/storage/managedwriter

     1  // Copyright 2022 Google LLC
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     https://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package managedwriter
    16  
    17  import (
    18  	"context"
    19  	"fmt"
    20  	"io"
    21  	"testing"
    22  
    23  	"github.com/googleapis/gax-go/v2/apierror"
    24  	"google.golang.org/grpc/codes"
    25  	"google.golang.org/grpc/status"
    26  )
    27  
    28  func TestManagedStream_AppendErrorRetries(t *testing.T) {
    29  
    30  	testCases := []struct {
    31  		err          error
    32  		attemptCount int
    33  		want         bool
    34  	}{
    35  		{
    36  			err:  nil,
    37  			want: false,
    38  		},
    39  		{
    40  			err:  fmt.Errorf("random error"),
    41  			want: false,
    42  		},
    43  		{
    44  			err:  io.EOF,
    45  			want: true,
    46  		},
    47  		{
    48  			err:          io.EOF,
    49  			attemptCount: 4,
    50  			want:         false,
    51  		},
    52  		{
    53  			err:  status.Error(codes.Unavailable, "nope"),
    54  			want: true,
    55  		},
    56  		{
    57  			err:  status.Error(codes.ResourceExhausted, "out of gas"),
    58  			want: false,
    59  		},
    60  		{
    61  			err:  status.Error(codes.ResourceExhausted, "Exceeds 'AppendRows throughput' quota for some reason"),
    62  			want: true,
    63  		},
    64  		{
    65  			err:  context.Canceled,
    66  			want: false,
    67  		},
    68  	}
    69  
    70  	retry := newStatelessRetryer()
    71  
    72  	for _, tc := range testCases {
    73  		if _, got := retry.Retry(tc.err, tc.attemptCount); got != tc.want {
    74  			t.Errorf("got %t, want %t for error: %+v", got, tc.want, tc.err)
    75  		}
    76  	}
    77  }
    78  
    79  func TestManagedStream_ShouldReconnect(t *testing.T) {
    80  
    81  	testCases := []struct {
    82  		err  error
    83  		want bool
    84  	}{
    85  		{
    86  			err:  fmt.Errorf("random error"),
    87  			want: false,
    88  		},
    89  		{
    90  			err:  io.EOF,
    91  			want: true,
    92  		},
    93  		{
    94  			err:  status.Error(codes.Unavailable, "the connection is draining"),
    95  			want: true,
    96  		},
    97  		{
    98  			err:  status.Error(codes.ResourceExhausted, "oof"), // may just be pushback
    99  			want: false,
   100  		},
   101  		{
   102  			err:  status.Error(codes.Canceled, "blah"),
   103  			want: true,
   104  		},
   105  		{
   106  			err:  status.Error(codes.Aborted, "connection has been idle too long"),
   107  			want: true,
   108  		},
   109  		{
   110  			err:  status.Error(codes.DeadlineExceeded, "blah"), // possibly bad backend, reconnect to speed recovery.
   111  			want: true,
   112  		},
   113  		{
   114  			err: func() error {
   115  				// wrap the underlying error in a gax apierror
   116  				ai, _ := apierror.FromError(status.Error(codes.Unavailable, "the connection is draining"))
   117  				return ai
   118  			}(),
   119  			want: true,
   120  		},
   121  	}
   122  
   123  	for _, tc := range testCases {
   124  		if got := shouldReconnect(tc.err); got != tc.want {
   125  			t.Errorf("got %t, want %t for error: %+v", got, tc.want, tc.err)
   126  		}
   127  	}
   128  }
   129
View as plain text