From 5c3131e8122b5b90d82bbe094ca923e60b15bddc Mon Sep 17 00:00:00 2001
From: Victor Schappert <schapper@amazon.com>
Date: Mon, 12 Dec 2022 16:33:56 -0800
Subject: [PATCH] Add more unit tests for retry on transient.UnexpectedEOF

---
 client_test.go | 62 ++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 62 insertions(+)

diff --git a/client_test.go b/client_test.go
index 3a1f8cb..f7a3166 100644
--- a/client_test.go
+++ b/client_test.go
@@ -627,6 +627,46 @@ func testClientRetryVarious(t *testing.T) {
 				assert.Nil(t, e.Body)
 			},
 		},
+		{
+			name:   "eof on do",
+			doResp: nil,
+			doErr: &url.Error{
+				Op:  "pling",
+				URL: "bling",
+				Err: io.EOF,
+			},
+			handlerCalls: []string{
+				"BeforeAttempt",
+				"AfterAttempt",
+			},
+			assertFunc: func(t *testing.T, e *request.Execution) {
+				require.IsType(t, &url.Error{}, e.Err)
+				urlError := e.Err.(*url.Error)
+				assert.Same(t, io.EOF, urlError.Err)
+				assert.Nil(t, e.Response)
+			},
+		},
+		{
+			name: "unexpected eof on read body",
+			doResp: &http.Response{
+				StatusCode: 200,
+				Body:       &unexpectedEOFReader{"I wish I could just...", 0},
+			},
+			doErr: nil,
+			handlerCalls: []string{
+				"BeforeAttempt",
+				"BeforeReadBody",
+				"AfterAttempt",
+			},
+			assertFunc: func(t *testing.T, e *request.Execution) {
+				require.IsType(t, &url.Error{}, e.Err)
+				urlError := e.Err.(*url.Error)
+				assert.Same(t, io.ErrUnexpectedEOF, urlError.Err)
+				assert.Equal(t, 200, e.StatusCode())
+				assert.NotNil(t, e.Response)
+				assert.Equal(t, []byte("I wish I could just..."), e.Body)
+			},
+		},
 		{
 			name: "no content",
 			doResp: &http.Response{
@@ -1819,3 +1859,25 @@ type callOrderMatcher struct {
 func (com *callOrderMatcher) Match(x int32) bool {
 	return atomic.CompareAndSwapInt32(&com.counter, x, x+1)
 }
+
+type unexpectedEOFReader struct {
+	s string
+	n int
+}
+
+func (r *unexpectedEOFReader) Read(p []byte) (n int, err error) {
+	n = len(p)
+	if n > len(r.s)-r.n {
+		n = len(r.s) - r.n
+	}
+	copy(p[0:n], []byte(r.s)[r.n:r.n+n])
+	r.n += n
+	if r.n == len(r.s) {
+		err = io.ErrUnexpectedEOF
+	}
+	return
+}
+
+func (r *unexpectedEOFReader) Close() error {
+	return nil
+}