Fix SDK's EndpointDiscovery and S3 Upload Manager unstable tests

aws/session/session.go

@@ -452,7 +452,7 @@ func mergeConfigSrcs(cfg, userCfg *aws.Config, envCfg envConfig, sharedCfg share
 		}
 	}
 
-	if aws.BoolValue(envCfg.EnableEndpointDiscovery) {
+	if cfg.EnableEndpointDiscovery == nil {
 		if envCfg.EnableEndpointDiscovery != nil {
 			cfg.WithEndpointDiscovery(*envCfg.EnableEndpointDiscovery)
 		} else if envCfg.EnableSharedConfig && sharedCfg.EnableEndpointDiscovery != nil {

private/model/api/codegentest/service/awsendpointdiscoverytest/endpoint_discovery_test.go

@@ -2,122 +2,97 @@ package awsendpointdiscoverytest
 
 import (
 	"sync"
+	"sync/atomic"
 	"testing"
 
 	"github.com/aws/aws-sdk-go/aws"
-	"github.com/aws/aws-sdk-go/aws/credentials"
 	"github.com/aws/aws-sdk-go/aws/request"
-	"github.com/aws/aws-sdk-go/aws/session"
+	"github.com/aws/aws-sdk-go/awstesting/unit"
 )
 
 func TestEndpointDiscovery(t *testing.T) {
-	cfg := &aws.Config{
-		Region:                  aws.String("mock-region"),
-		Credentials:             credentials.NewStaticCredentials("AKID", "SECRET", "SESSION"),
+	svc := New(unit.Session, &aws.Config{
 		EnableEndpointDiscovery: aws.Bool(true),
-	}
-	sess := session.New(cfg)
-	sess.Handlers = removeHandlers(sess.Handlers, true)
+	})
+	svc.Handlers.Clear()
+	svc.Handlers.Send.PushBack(mockSendDescEndpoint)
 
-	sess.Handlers.Send.PushBack(func(r *request.Request) {
-		out, ok := r.Data.(*DescribeEndpointsOutput)
-		if !ok {
+	var descCount int32
+	svc.Handlers.Complete.PushBack(func(r *request.Request) {
+		if r.Operation.Name != opDescribeEndpoints {
 			return
 		}
-
-		out.Endpoints = []*Endpoint{
-			{
-				Address:              aws.String("http://foo"),
-				CachePeriodInMinutes: aws.Int64(5),
-			},
-		}
-		r.Data = out
-	})
-
-	svc := New(sess)
-	svc.Handlers = removeHandlers(svc.Handlers, false)
-
-	req, _ := svc.TestDiscoveryIdentifiersRequiredRequest(&TestDiscoveryIdentifiersRequiredInput{
-		Sdk: aws.String("sdk"),
+		atomic.AddInt32(&descCount, 1)
 	})
 
-	req.Handlers = removeHandlers(req.Handlers, false)
-
-	req.Handlers.Send.PushBack(func(r *request.Request) {
-		if e, a := "http://foo", r.HTTPRequest.URL.String(); e != a {
-			t.Errorf("expected %q, but received %q", e, a)
+	for i := 0; i < 2; i++ {
+		req, _ := svc.TestDiscoveryIdentifiersRequiredRequest(
+			&TestDiscoveryIdentifiersRequiredInput{
+				Sdk: aws.String("sdk"),
+			},
+		)
+		req.Handlers.Send.PushBack(func(r *request.Request) {
+			if e, a := "http://foo", r.HTTPRequest.URL.String(); e != a {
+				t.Errorf("expected %q, but received %q", e, a)
+			}
+		})
+		if err := req.Send(); err != nil {
+			t.Fatal(err)
 		}
-	})
+	}
 
-	if err := req.Send(); err != nil {
-		t.Fatal(err)
+	if e, a := int32(1), atomic.LoadInt32(&descCount); e != a {
+		t.Errorf("expect desc endpoint called %d, got %d", e, a)
 	}
 }
 
 func TestAsyncEndpointDiscovery(t *testing.T) {
-	cfg := &aws.Config{
-		Region:                  aws.String("mock-region"),
-		Credentials:             credentials.NewStaticCredentials("AKID", "SECRET", "SESSION"),
+	svc := New(unit.Session, &aws.Config{
 		EnableEndpointDiscovery: aws.Bool(true),
-	}
-	sess := session.New(cfg)
-	sess.Handlers = removeHandlers(sess.Handlers, true)
+	})
+	svc.Handlers.Clear()
 
-	sess.Handlers.Send.PushBack(func(r *request.Request) {
-		out, ok := r.Data.(*DescribeEndpointsOutput)
-		if !ok {
-			return
+	var firstAsyncReq sync.WaitGroup
+	firstAsyncReq.Add(1)
+	svc.Handlers.Send.PushBack(func(r *request.Request) {
+		if r.Operation.Name == opDescribeEndpoints {
+			firstAsyncReq.Wait()
 		}
-
-		out.Endpoints = []*Endpoint{
-			{
-				Address:              aws.String("http://foo"),
-				CachePeriodInMinutes: aws.Int64(5),
-			},
-		}
-		r.Data = out
 	})
+	svc.Handlers.Send.PushBack(mockSendDescEndpoint)
 
-	svc := New(sess)
-	svc.Handlers = removeHandlers(svc.Handlers, false)
-
-	var wg sync.WaitGroup
-	wg.Add(1)
+	var descWg sync.WaitGroup
+	descWg.Add(1)
 	svc.Handlers.Complete.PushBack(func(r *request.Request) {
-		if r.Operation.Name == "DescribeEndpoints" {
-			wg.Done()
+		if r.Operation.Name == opDescribeEndpoints {
+			descWg.Done()
 		}
 	})
 
 	req, _ := svc.TestDiscoveryOptionalRequest(&TestDiscoveryOptionalInput{
 		Sdk: aws.String("sdk"),
 	})
-
-	req.Handlers = removeHandlers(req.Handlers, false)
-
-	const expectedURI = "awsendpointdiscoverytestservice.mock-region.amazonaws.com"
+	const clientHost = "awsendpointdiscoverytestservice.mock-region.amazonaws.com"
 	req.Handlers.Send.PushBack(func(r *request.Request) {
-		if e, a := expectedURI, r.HTTPRequest.URL.Host; e != a {
+		if e, a := clientHost, r.HTTPRequest.URL.Host; e != a {
 			t.Errorf("expected %q, but received %q", e, a)
 		}
 	})
-
 	if err := req.Send(); err != nil {
 		t.Fatal(err)
 	}
 
-	wg.Wait()
+	firstAsyncReq.Done()
+	descWg.Wait()
+
 	req, _ = svc.TestDiscoveryOptionalRequest(&TestDiscoveryOptionalInput{
 		Sdk: aws.String("sdk"),
 	})
-
-	req.Handlers = removeHandlers(req.Handlers, false)
 	req.Handlers.Send.PushBack(func(r *request.Request) {
 		if e, a := "http://foo", r.HTTPRequest.URL.String(); e != a {
 			t.Errorf("expected %q, but received %q", e, a)
 		}
 	})
-
 	if err := req.Send(); err != nil {
 		t.Fatal(err)
 	}
@@ -136,3 +111,18 @@ func removeHandlers(h request.Handlers, removeSendHandlers bool) request.Handler
 	h.ValidateResponse.Clear()
 	return h
 }
+
+func mockSendDescEndpoint(r *request.Request) {
+	if r.Operation.Name != opDescribeEndpoints {
+		return
+	}
+
+	out, _ := r.Data.(*DescribeEndpointsOutput)
+	out.Endpoints = []*Endpoint{
+		{
+			Address:              aws.String("http://foo"),
+			CachePeriodInMinutes: aws.Int64(5),
+		},
+	}
+	r.Data = out
+}

service/s3/s3manager/upload_test.go

@@ -785,53 +785,45 @@ func TestUploadInputS3PutObjectInputPairity(t *testing.T) {
 }
 
 type testIncompleteReader struct {
-	Buf   []byte
-	Count int
+	Size int64
+	read int64
 }
 
 func (r *testIncompleteReader) Read(p []byte) (n int, err error) {
-	if r.Count < 0 {
-		return 0, io.ErrUnexpectedEOF
+	r.read += int64(len(p))
+	if r.read >= r.Size {
+		return int(r.read - r.Size), io.ErrUnexpectedEOF
 	}
-
-	r.Count--
-	return copy(p, r.Buf), nil
+	return len(p), nil
 }
 
 func TestUploadUnexpectedEOF(t *testing.T) {
-	s, ops, args := loggingSvc(emptyList)
+	s, ops, _ := loggingSvc(emptyList)
 	mgr := s3manager.NewUploaderWithClient(s, func(u *s3manager.Uploader) {
 		u.Concurrency = 1
+		u.PartSize = s3manager.MinUploadPartSize
 	})
 	_, err := mgr.Upload(&s3manager.UploadInput{
 		Bucket: aws.String("Bucket"),
 		Key:    aws.String("Key"),
 		Body: &testIncompleteReader{
-			Buf:   make([]byte, 1024*1024*5),
-			Count: 1,
+			Size: int64(s3manager.MinUploadPartSize + 1),
 		},
 	})
-
 	if err == nil {
 		t.Error("Expected error, but received none")
 	}
 
+	// Ensure upload started.
 	if e, a := "CreateMultipartUpload", (*ops)[0]; e != a {
 		t.Errorf("Expected %q, but received %q", e, a)
 	}
 
-	if e, a := "UploadPart", (*ops)[1]; e != a {
-		t.Errorf("Expected %q, but received %q", e, a)
-	}
-
+	// Part may or may not be sent because of timing of sending parts and
+	// reading next part in upload manager. Just check for the last abort.
 	if e, a := "AbortMultipartUpload", (*ops)[len(*ops)-1]; e != a {
 		t.Errorf("Expected %q, but received %q", e, a)
 	}
-
-	// Part lengths
-	if e, a := 1024*1024*5, buflen(val((*args)[1], "Body")); e != a {
-		t.Errorf("Expected %d, but received %d", e, a)
-	}
 }
 
 func compareStructType(a, b reflect.Type) map[string]int {