Fix mocks

Signed-off-by: Kazuyoshi Kato <[email protected]>
aws · Mar 15, 2023 · 621a4b1 · 621a4b1
1 parent 7a0e6d1
commit 621a4b1
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 5 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/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),