From 621a4b169670d94e527073a3d64a8ed68cc7772a Mon Sep 17 00:00:00 2001 From: Kazuyoshi Kato Date: Wed, 15 Mar 2023 16:57:16 -0700 Subject: [PATCH] Fix mocks Signed-off-by: Kazuyoshi Kato --- agent/config/config_test.go | 1 - agent/s3/s3_test.go | 6 ++++-- agent/taskresource/envFiles/envfile_test.go | 7 +++++-- 3 files changed, 9 insertions(+), 5 deletions(-) diff --git a/agent/config/config_test.go b/agent/config/config_test.go index 9c563ff8c0e..f8fabf4705f 100644 --- 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 index 6186255c788..9610d58cd3b 100644 --- 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 index ea94c32da11..ae0af291615 100644 --- 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),