Fix mocks

mockgen 1.6 is more strict than 1.3. Signed-off-by: Kazuyoshi Kato <[email protected]>
aws · Mar 16, 2023 · 876d626 · 876d626
1 parent 7a0e6d1
commit 876d626
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 13 deletions.
diff --git a/agent/config/config_test.go b/agent/config/config_test.go
@@ -88,7 +88,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
@@ -1414,9 +1414,6 @@ func TestTaskWaitForExecutionCredentials(t *testing.T) {
 
 	for _, tc := range tcs {
 		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{
@@ -1425,12 +1422,9 @@ func TestTaskWaitForExecutionCredentials(t *testing.T) {
 					KnownStatusUnsafe:   apitaskstatus.TaskRunning,
 					DesiredStatusUnsafe: apitaskstatus.TaskRunning,
 				},
-				_time:       mockTime,
 				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
@@ -25,6 +25,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"
 
@@ -46,8 +47,9 @@ func TestDownloadFile(t *testing.T) {
 	mockFile := mock_oswrapper.NewMockFile()
 	mockS3ManagerClient := mock_s3manager.NewMockS3ManagerClient(ctrl)
 
-	mockS3ManagerClient.EXPECT().DownloadWithContext(gomock.Any(), mockFile, gomock.Any()).Do(func(ctx aws.Context,
-		w io.WriterAt, input *s3sdk.GetObjectInput) {
+	mockS3ManagerClient.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))
 	})

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,10 @@ func TestCreateWithEnvVarFile(t *testing.T) {
 		mockCredentialsManager.EXPECT().GetTaskCredentials(executionCredentialsID).Return(creds, true),
 		mockS3ClientCreator.EXPECT().NewS3ManagerClient(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),

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().NewS3ManagerClient("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),