Fix mocks

gomock.Controller automatically calls Finish() which finds some broken mocks. Signed-off-by: Kazuyoshi Kato <[email protected]>
aws · Nov 2, 2022 · abe0dbb · abe0dbb
1 parent 64ace3b
commit abe0dbb
Show file tree

Hide file tree

Showing 6 changed files with 18 additions and 19 deletions.
diff --git a/agent/config/config_test.go b/agent/config/config_test.go
@@ -86,7 +86,6 @@ func TestBrokenEC2MetadataEndpoint(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	mockEc2Metadata := mock_ec2.NewMockEC2MetadataClient(ctrl)
 
-	mockEc2Metadata.EXPECT().InstanceIdentityDocument().Return(ec2metadata.EC2InstanceIdentityDocument{}, errors.New("err"))
 	mockEc2Metadata.EXPECT().GetUserData()
 
 	config, err := NewConfig(mockEc2Metadata)

diff --git a/agent/engine/task_manager_test.go b/agent/engine/task_manager_test.go
@@ -1416,7 +1416,6 @@ func TestTaskWaitForExecutionCredentials(t *testing.T) {
 		t.Run(fmt.Sprintf("%v", tc.errs), func(t *testing.T) {
 			ctrl := gomock.NewController(t)
 			mockTime := mock_ttime.NewMockTime(ctrl)
-			mockTimer := mock_ttime.NewMockTimer(ctrl)
 			ctx, cancel := context.WithCancel(context.TODO())
 			defer cancel()
 			task := &managedTask{
@@ -1429,8 +1428,6 @@ func TestTaskWaitForExecutionCredentials(t *testing.T) {
 				acsMessages: make(chan acsTransition),
 			}
 			if tc.result {
-				mockTime.EXPECT().AfterFunc(gomock.Any(), gomock.Any()).Return(mockTimer)
-				mockTimer.EXPECT().Stop()
 				go func() { task.acsMessages <- acsTransition{desiredStatus: apitaskstatus.TaskRunning} }()
 			}
 

diff --git a/agent/s3/s3_test.go b/agent/s3/s3_test.go
@@ -24,6 +24,7 @@ import (
 
 	"github.com/aws/aws-sdk-go/aws"
 	s3sdk "github.com/aws/aws-sdk-go/service/s3"
+	"github.com/aws/aws-sdk-go/service/s3/s3manager"
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
 
@@ -44,11 +45,11 @@ func TestDownloadFile(t *testing.T) {
 	mockFile := mock_oswrapper.NewMockFile()
 	mockS3Client := mock_s3.NewMockS3Client(ctrl)
 
-	mockS3Client.EXPECT().DownloadWithContext(gomock.Any(), mockFile, gomock.Any()).Do(func(ctx aws.Context,
-		w io.WriterAt, input *s3sdk.GetObjectInput) {
-		assert.Equal(t, testBucket, aws.StringValue(input.Bucket))
-		assert.Equal(t, testKey, aws.StringValue(input.Key))
-	})
+	mockS3Client.EXPECT().DownloadWithContext(gomock.Any(), mockFile, gomock.Any(), gomock.Any()).
+		Do(func(ctx aws.Context, w io.WriterAt, input *s3sdk.GetObjectInput, options ...func(*s3manager.Downloader)) {
+			assert.Equal(t, testBucket, aws.StringValue(input.Bucket))
+			assert.Equal(t, testKey, aws.StringValue(input.Key))
+		})
 
 	err := DownloadFile(testBucket, testKey, testTimeout, mockFile, mockS3Client)
 	assert.NoError(t, err)
@@ -61,7 +62,7 @@ func TestDownloadFileError(t *testing.T) {
 	mockS3Client := mock_s3.NewMockS3Client(ctrl)
 	mockFile := mock_oswrapper.NewMockFile()
 
-	mockS3Client.EXPECT().DownloadWithContext(gomock.Any(), mockFile, gomock.Any()).Return(int64(0), errors.New("test error"))
+	mockS3Client.EXPECT().DownloadWithContext(gomock.Any(), mockFile, gomock.Any(), gomock.Any()).Return(int64(0), errors.New("test error"))
 
 	err := DownloadFile(testBucket, testKey, testTimeout, mockFile, mockS3Client)
 	assert.Error(t, err)

diff --git a/agent/taskresource/credentialspec/credentialspec_windows_test.go b/agent/taskresource/credentialspec/credentialspec_windows_test.go
@@ -497,7 +497,7 @@ func TestHandleS3CredentialspecFile(t *testing.T) {
 	gomock.InOrder(
 		s3ClientCreator.EXPECT().NewS3ClientForBucket(gomock.Any(), gomock.Any(), gomock.Any()).Return(mockS3Client, nil),
 		mockIO.EXPECT().TempFile(gomock.Any(), gomock.Any()).Return(mockFile, nil),
-		mockS3Client.EXPECT().DownloadWithContext(gomock.Any(), mockFile, gomock.Any()).Return(int64(0), nil),
+		mockS3Client.EXPECT().DownloadWithContext(gomock.Any(), mockFile, gomock.Any(), gomock.Any()).Return(int64(0), nil),
 	)
 
 	err := cs.handleS3CredentialspecFile(s3CredentialSpec, credentialSpecS3ARN, iamCredentials)
@@ -607,7 +607,7 @@ func TestHandleS3CredentialspecFileWriteErr(t *testing.T) {
 	gomock.InOrder(
 		s3ClientCreator.EXPECT().NewS3ClientForBucket(gomock.Any(), gomock.Any(), gomock.Any()).Return(mockS3Client, nil),
 		mockIO.EXPECT().TempFile(gomock.Any(), gomock.Any()).Return(mockFile, nil),
-		mockS3Client.EXPECT().DownloadWithContext(gomock.Any(), mockFile, gomock.Any()).Return(int64(0), nil),
+		mockS3Client.EXPECT().DownloadWithContext(gomock.Any(), mockFile, gomock.Any(), gomock.Any()).Return(int64(0), nil),
 	)
 
 	err := cs.handleS3CredentialspecFile(s3CredentialSpec, credentialSpecS3ARN, iamCredentials)
@@ -727,7 +727,7 @@ func TestCreateS3(t *testing.T) {
 		credentialsManager.EXPECT().GetTaskCredentials(gomock.Any()).Return(creds, true),
 		s3ClientCreator.EXPECT().NewS3ClientForBucket(gomock.Any(), gomock.Any(), gomock.Any()).Return(mockS3Client, nil),
 		mockIO.EXPECT().TempFile(gomock.Any(), gomock.Any()).Return(mockFile, nil),
-		mockS3Client.EXPECT().DownloadWithContext(gomock.Any(), mockFile, gomock.Any()).Return(int64(0), nil),
+		mockS3Client.EXPECT().DownloadWithContext(gomock.Any(), mockFile, gomock.Any(), gomock.Any()).Return(int64(0), nil),
 	)
 
 	assert.NoError(t, cs.Create())

diff --git a/agent/taskresource/envFiles/envfile_test.go b/agent/taskresource/envFiles/envfile_test.go
@@ -37,6 +37,7 @@ import (
 	mock_oswrapper "github.com/aws/amazon-ecs-agent/agent/utils/oswrapper/mocks"
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/s3"
+	"github.com/aws/aws-sdk-go/service/s3/s3manager"
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
 )
@@ -141,8 +142,8 @@ func TestCreateWithEnvVarFile(t *testing.T) {
 		mockCredentialsManager.EXPECT().GetTaskCredentials(executionCredentialsID).Return(creds, true),
 		mockS3ClientCreator.EXPECT().NewS3ClientForBucket(s3Bucket, region, creds.IAMRoleCredentials).Return(mockS3Client, nil),
 		mockIOUtil.EXPECT().TempFile(resourceDir, gomock.Any()).Return(mockFile, nil),
-		mockS3Client.EXPECT().DownloadWithContext(gomock.Any(), mockFile, gomock.Any()).Do(
-			func(ctx aws.Context, w io.WriterAt, input *s3.GetObjectInput) {
+		mockS3Client.EXPECT().DownloadWithContext(gomock.Any(), mockFile, gomock.Any(), gomock.Any()).Do(
+			func(ctx aws.Context, w io.WriterAt, input *s3.GetObjectInput, options ...func(*s3manager.Downloader)) {
 				assert.Equal(t, s3Bucket, aws.StringValue(input.Bucket))
 				assert.Equal(t, s3Key, aws.StringValue(input.Key))
 			}).Return(int64(0), nil),
@@ -258,7 +259,7 @@ func TestCreateRenameFileError(t *testing.T) {
 		mockCredentialsManager.EXPECT().GetTaskCredentials(executionCredentialsID).Return(creds, true),
 		mockS3ClientCreator.EXPECT().NewS3ClientForBucket(s3Bucket, region, creds.IAMRoleCredentials).Return(mockS3Client, nil),
 		mockIOUtil.EXPECT().TempFile(resourceDir, gomock.Any()).Return(mockFile, nil),
-		mockS3Client.EXPECT().DownloadWithContext(gomock.Any(), mockFile, gomock.Any()).Return(int64(0), nil),
+		mockS3Client.EXPECT().DownloadWithContext(gomock.Any(), mockFile, gomock.Any(), gomock.Any()).Return(int64(0), nil),
 	)
 
 	assert.Error(t, envfileResource.Create())

diff --git a/agent/taskresource/firelens/firelens_unix_test.go b/agent/taskresource/firelens/firelens_unix_test.go
@@ -23,6 +23,7 @@ import (
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/s3"
+	"github.com/aws/aws-sdk-go/service/s3/s3manager"
 	"github.com/golang/mock/gomock"
 	"github.com/pkg/errors"
 	"github.com/stretchr/testify/assert"
@@ -366,8 +367,8 @@ func TestCreateFirelensResourceWithS3Config(t *testing.T) {
 		mockS3ClientCreator.EXPECT().NewS3ClientForBucket("bucket", testRegion, creds.IAMRoleCredentials).Return(mockS3Client, nil),
 		// write external config file downloaded from s3
 		mockIOUtil.EXPECT().TempFile(testResourceDir, tempFile).Return(mockFile, nil),
-		mockS3Client.EXPECT().DownloadWithContext(gomock.Any(), mockFile, gomock.Any()).Do(
-			func(ctx aws.Context, w io.WriterAt, input *s3.GetObjectInput) {
+		mockS3Client.EXPECT().DownloadWithContext(gomock.Any(), mockFile, gomock.Any(), gomock.Any()).Do(
+			func(ctx aws.Context, w io.WriterAt, input *s3.GetObjectInput, options ...func(*s3manager.Downloader)) {
 				assert.Equal(t, "bucket", aws.StringValue(input.Bucket))
 				assert.Equal(t, "key", aws.StringValue(input.Key))
 			}).Return(int64(0), nil),
@@ -437,7 +438,7 @@ func TestCreateFirelensResourceWithS3ConfigDownloadFailure(t *testing.T) {
 		mockCredentialsManager.EXPECT().GetTaskCredentials(testExecutionCredentialsID).Return(creds, true),
 		mockS3ClientCreator.EXPECT().NewS3ClientForBucket("bucket", testRegion, creds.IAMRoleCredentials).Return(mockS3Client, nil),
 		mockIOUtil.EXPECT().TempFile(testResourceDir, tempFile).Return(mockFile, nil),
-		mockS3Client.EXPECT().DownloadWithContext(gomock.Any(), mockFile, gomock.Any()).Return(int64(0), errors.New("test error")),
+		mockS3Client.EXPECT().DownloadWithContext(gomock.Any(), mockFile, gomock.Any(), gomock.Any()).Return(int64(0), errors.New("test error")),
 	)
 
 	assert.Error(t, firelensResource.Create())