github.com/aavshr/aws-sdk-go@v1.41.3/service/kinesis/customizations_test.go (about) 1 package kinesis 2 3 import ( 4 "bytes" 5 "fmt" 6 "io" 7 "io/ioutil" 8 "net/http" 9 "testing" 10 "time" 11 12 "github.com/aavshr/aws-sdk-go/aws" 13 "github.com/aavshr/aws-sdk-go/aws/awserr" 14 "github.com/aavshr/aws-sdk-go/aws/corehandlers" 15 "github.com/aavshr/aws-sdk-go/aws/request" 16 "github.com/aavshr/aws-sdk-go/awstesting/unit" 17 ) 18 19 type testReader struct { 20 duration time.Duration 21 } 22 23 func (r *testReader) Read(b []byte) (int, error) { 24 time.Sleep(r.duration) 25 return 0, io.EOF 26 } 27 28 func (r *testReader) Close() error { 29 return nil 30 } 31 32 // GetRecords will hang unexpectedly during reads. 33 // See https://github.com/aavshr/aws-sdk-go/issues/1141 34 func TestKinesisGetRecordsCustomization(t *testing.T) { 35 readDuration = time.Millisecond 36 retryCount := 0 37 svc := New(unit.Session, &aws.Config{ 38 MaxRetries: aws.Int(4), 39 }) 40 req, _ := svc.GetRecordsRequest(&GetRecordsInput{ 41 ShardIterator: aws.String("foo"), 42 }) 43 req.Handlers.Send.Clear() 44 req.Handlers.Send.PushBack(func(r *request.Request) { 45 r.HTTPResponse = &http.Response{ 46 StatusCode: 200, 47 Header: http.Header{ 48 "X-Amz-Request-Id": []string{"abc123"}, 49 }, 50 Body: &testReader{duration: 10 * time.Second}, 51 ContentLength: -1, 52 } 53 r.HTTPResponse.Status = http.StatusText(r.HTTPResponse.StatusCode) 54 retryCount++ 55 }) 56 req.ApplyOptions(request.WithResponseReadTimeout(time.Second)) 57 err := req.Send() 58 if err == nil { 59 t.Errorf("Expected error, but received nil") 60 } else if v, ok := err.(awserr.Error); !ok { 61 t.Errorf("Expected awserr.Error but received %v", err) 62 } else if v.Code() != request.ErrCodeResponseTimeout { 63 t.Errorf("Expected 'RequestTimeout' error, but received %s instead", v.Code()) 64 } 65 if retryCount != 5 { 66 t.Errorf("Expected '5' retries, but received %d", retryCount) 67 } 68 } 69 70 func TestKinesisGetRecordsNoTimeout(t *testing.T) { 71 readDuration = time.Second 72 svc := New(unit.Session) 73 req, _ := svc.GetRecordsRequest(&GetRecordsInput{ 74 ShardIterator: aws.String("foo"), 75 }) 76 req.Handlers.Send.Clear() 77 req.Handlers.Send.PushBack(func(r *request.Request) { 78 r.HTTPResponse = &http.Response{ 79 StatusCode: 200, 80 Header: http.Header{ 81 "X-Amz-Request-Id": []string{"abc123"}, 82 }, 83 Body: &testReader{duration: time.Duration(0)}, 84 ContentLength: -1, 85 } 86 r.HTTPResponse.Status = http.StatusText(r.HTTPResponse.StatusCode) 87 }) 88 req.ApplyOptions(request.WithResponseReadTimeout(time.Second)) 89 err := req.Send() 90 if err != nil { 91 t.Errorf("Expected no error, but received %v", err) 92 } 93 } 94 95 func TestKinesisCustomRetryErrorCodes(t *testing.T) { 96 svc := New(unit.Session, &aws.Config{ 97 MaxRetries: aws.Int(1), 98 LogLevel: aws.LogLevel(aws.LogDebugWithHTTPBody), 99 }) 100 svc.Handlers.Validate.Clear() 101 102 const jsonErr = `{"__type":%q, "message":"some error message"}` 103 var reqCount int 104 resps := []*http.Response{ 105 { 106 StatusCode: 400, 107 Header: http.Header{}, 108 Body: ioutil.NopCloser(bytes.NewReader( 109 []byte(fmt.Sprintf(jsonErr, ErrCodeLimitExceededException)), 110 )), 111 }, 112 { 113 StatusCode: 200, 114 Header: http.Header{}, 115 Body: ioutil.NopCloser(bytes.NewReader([]byte{})), 116 }, 117 } 118 119 req, _ := svc.GetRecordsRequest(&GetRecordsInput{}) 120 req.Handlers.Send.Swap(corehandlers.SendHandler.Name, request.NamedHandler{ 121 Name: "custom send handler", 122 Fn: func(r *request.Request) { 123 r.HTTPResponse = resps[reqCount] 124 reqCount++ 125 }, 126 }) 127 128 if err := req.Send(); err != nil { 129 t.Fatalf("expect no error, got %v", err) 130 } 131 132 if e, a := 2, reqCount; e != a { 133 t.Errorf("expect %v requests, got %v", e, a) 134 } 135 }