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  }