diff --git a/provider/go.mod b/provider/go.mod index e76974c8f9..3c849f68a3 100644 --- a/provider/go.mod +++ b/provider/go.mod @@ -4,16 +4,19 @@ go 1.21 require ( github.com/apparentlymart/go-cidr v1.1.0 + github.com/aws/aws-lambda-go v1.47.0 github.com/aws/aws-sdk-go v1.50.36 - github.com/aws/aws-sdk-go-v2 v1.26.1 + github.com/aws/aws-sdk-go-v2 v1.32.3 github.com/aws/aws-sdk-go-v2/config v1.27.11 github.com/aws/aws-sdk-go-v2/credentials v1.17.11 github.com/aws/aws-sdk-go-v2/service/cloudcontrol v1.15.7 github.com/aws/aws-sdk-go-v2/service/cloudformation v1.43.0 github.com/aws/aws-sdk-go-v2/service/ec2 v1.146.0 + github.com/aws/aws-sdk-go-v2/service/lambda v1.64.1 + github.com/aws/aws-sdk-go-v2/service/s3 v1.66.2 github.com/aws/aws-sdk-go-v2/service/ssm v1.49.2 github.com/aws/aws-sdk-go-v2/service/sts v1.28.6 - github.com/aws/smithy-go v1.20.2 + github.com/aws/smithy-go v1.22.0 github.com/blang/semver v3.5.1+incompatible github.com/goccy/go-yaml v1.9.5 github.com/golang/glog v1.2.0 @@ -58,12 +61,16 @@ require ( github.com/agext/levenshtein v1.2.3 // indirect github.com/apparentlymart/go-textseg/v13 v13.0.0 // indirect github.com/atotto/clipboard v0.1.4 // indirect + github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.6 // indirect github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.1 // indirect - github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.5 // indirect - github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.5 // indirect + github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.22 // indirect + github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.22 // indirect github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0 // indirect - github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.2 // indirect - github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.7 // indirect + github.com/aws/aws-sdk-go-v2/internal/v4a v1.3.22 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.0 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.4.3 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.3 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.18.3 // indirect github.com/aws/aws-sdk-go-v2/service/kms v1.30.1 // indirect github.com/aws/aws-sdk-go-v2/service/sso v1.20.5 // indirect github.com/aws/aws-sdk-go-v2/service/ssooidc v1.23.4 // indirect diff --git a/provider/go.sum b/provider/go.sum index 24b5ad0214..74a4366291 100644 --- a/provider/go.sum +++ b/provider/go.sum @@ -91,12 +91,14 @@ github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5 h1:0CwZNZbxp69SHPd github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5/go.mod h1:wHh0iHkYZB8zMSxRWpUBQtwG5a7fFgvEO+odwuTv2gs= github.com/atotto/clipboard v0.1.4 h1:EH0zSVneZPSuFR11BlR9YppQTVDbh5+16AmcJi4g1z4= github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI= +github.com/aws/aws-lambda-go v1.47.0 h1:0H8s0vumYx/YKs4sE7YM0ktwL2eWse+kfopsRI1sXVI= +github.com/aws/aws-lambda-go v1.47.0/go.mod h1:dpMpZgvWx5vuQJfBt0zqBha60q7Dd7RfgJv23DymV8A= github.com/aws/aws-sdk-go v1.50.36 h1:PjWXHwZPuTLMR1NIb8nEjLucZBMzmf84TLoLbD8BZqk= github.com/aws/aws-sdk-go v1.50.36/go.mod h1:LF8svs817+Nz+DmiMQKTO3ubZ/6IaTpq3TjupRn3Eqk= -github.com/aws/aws-sdk-go-v2 v1.26.1 h1:5554eUqIYVWpU0YmeeYZ0wU64H2VLBs8TlhRB2L+EkA= -github.com/aws/aws-sdk-go-v2 v1.26.1/go.mod h1:ffIFB97e2yNsv4aTSGkqtHnppsIJzw7G7BReUZ3jCXM= -github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.2 h1:x6xsQXGSmW6frevwDA+vi/wqhp1ct18mVXYN08/93to= -github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.2/go.mod h1:lPprDr1e6cJdyYeGXnRaJoP4Md+cDBvi2eOj00BlGmg= +github.com/aws/aws-sdk-go-v2 v1.32.3 h1:T0dRlFBKcdaUPGNtkBSwHZxrtis8CQU17UpNBZYd0wk= +github.com/aws/aws-sdk-go-v2 v1.32.3/go.mod h1:2SK5n0a2karNTv5tbP1SjsX0uhttou00v/HpXKM1ZUo= +github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.6 h1:pT3hpW0cOHRJx8Y0DfJUEQuqPild8jRGmSFmBgvydr0= +github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.6/go.mod h1:j/I2++U0xX+cr44QjHay4Cvxj6FUbnxrgmqN3H1jTZA= github.com/aws/aws-sdk-go-v2/config v1.27.11 h1:f47rANd2LQEYHda2ddSCKYId18/8BhSRM4BULGmfgNA= github.com/aws/aws-sdk-go-v2/config v1.27.11/go.mod h1:SMsV78RIOYdve1vf36z8LmnszlRWkwMQtomCAI0/mIE= github.com/aws/aws-sdk-go-v2/credentials v1.17.11 h1:YuIB1dJNf1Re822rriUOTxopaHHvIq0l/pX3fwO+Tzs= @@ -105,14 +107,14 @@ github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.1 h1:FVJ0r5XTHSmIHJV6KuDmdYh github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.1/go.mod h1:zusuAeqezXzAB24LGuzuekqMAEgWkVYukBec3kr3jUg= github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.16.15 h1:7Zwtt/lP3KNRkeZre7soMELMGNoBrutx8nobg1jKWmo= github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.16.15/go.mod h1:436h2adoHb57yd+8W+gYPrrA9U/R/SuAuOO42Ushzhw= -github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.5 h1:aw39xVGeRWlWx9EzGVnhOR4yOjQDHPQ6o6NmBlscyQg= -github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.5/go.mod h1:FSaRudD0dXiMPK2UjknVwwTYyZMRsHv3TtkabsZih5I= -github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.5 h1:PG1F3OD1szkuQPzDw3CIQsRIrtTlUC3lP84taWzHlq0= -github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.5/go.mod h1:jU1li6RFryMz+so64PpKtudI+QzbKoIEivqdf6LNpOc= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.22 h1:Jw50LwEkVjuVzE1NzkhNKkBf9cRN7MtE1F/b2cOKTUM= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.22/go.mod h1:Y/SmAyPcOTmpeVaWSzSKiILfXTVJwrGmYZhcRbhWuEY= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.22 h1:981MHwBaRZM7+9QSR6XamDzF/o7ouUGxFzr+nVSIhrs= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.22/go.mod h1:1RA1+aBEfn+CAB/Mh0MB6LsdCYCnjZm7tKXtnk499ZQ= github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0 h1:hT8rVHwugYE2lEfdFE0QWVo81lF7jMrYJVDWI+f+VxU= github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0/go.mod h1:8tu/lYfQfFe6IGnaOdrpVgEL2IrrDOf6/m9RQum4NkY= -github.com/aws/aws-sdk-go-v2/internal/v4a v1.3.5 h1:81KE7vaZzrl7yHBYHVEzYB8sypz11NMOZ40YlWvPxsU= -github.com/aws/aws-sdk-go-v2/internal/v4a v1.3.5/go.mod h1:LIt2rg7Mcgn09Ygbdh/RdIm0rQ+3BNkbP1gyVMFtRK0= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.3.22 h1:yV+hCAHZZYJQcwAaszoBNwLbPItHvApxT0kVIw6jRgs= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.3.22/go.mod h1:kbR1TL8llqB1eGnVbybcA4/wgScxdylOdyAd51yxPdw= github.com/aws/aws-sdk-go-v2/service/cloudcontrol v1.15.7 h1:8sBfx7QkDZ6dgfUNXWHWRc6Eax7WOI3Slgj6OKDHKTI= github.com/aws/aws-sdk-go-v2/service/cloudcontrol v1.15.7/go.mod h1:P1EMD13hrBE2KUw030w482Eyk2NmOFIvGqmgNi4XRDc= github.com/aws/aws-sdk-go-v2/service/cloudformation v1.43.0 h1:fusTelL7ZIvR51E+xwc/HVUlWGhkWFlS+dtYrynVBq4= @@ -121,18 +123,20 @@ github.com/aws/aws-sdk-go-v2/service/ec2 v1.146.0 h1:d6pYx/CKADORpxqBINY7DuD4V1f github.com/aws/aws-sdk-go-v2/service/ec2 v1.146.0/go.mod h1:hIsHE0PaWAQakLCshKS7VKWMGXaqrAFp4m95s2W9E6c= github.com/aws/aws-sdk-go-v2/service/iam v1.31.4 h1:eVm30ZIDv//r6Aogat9I88b5YX1xASSLcEDqHYRPVl0= github.com/aws/aws-sdk-go-v2/service/iam v1.31.4/go.mod h1:aXWImQV0uTW35LM0A/T4wEg6R1/ReXUu4SM6/lUHYK0= -github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.2 h1:Ji0DY1xUsUr3I8cHps0G+XM3WWU16lP6yG8qu1GAZAs= -github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.2/go.mod h1:5CsjAbs3NlGQyZNFACh+zztPDI7fU6eW9QsxjfnuBKg= -github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.3.7 h1:ZMeFZ5yk+Ek+jNr1+uwCd2tG89t6oTS5yVWpa6yy2es= -github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.3.7/go.mod h1:mxV05U+4JiHqIpGqqYXOHLPKUC6bDXC44bsUhNjOEwY= -github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.7 h1:ogRAwT1/gxJBcSWDMZlgyFUM962F51A5CRhDLbxLdmo= -github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.7/go.mod h1:YCsIZhXfRPLFFCl5xxY+1T9RKzOKjCut+28JSX2DnAk= -github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.17.5 h1:f9RyWNtS8oH7cZlbn+/JNPpjUk5+5fLd5lM9M0i49Ys= -github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.17.5/go.mod h1:h5CoMZV2VF297/VLhRhO1WF+XYWOzXo+4HsObA4HjBQ= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.0 h1:TToQNkvGguu209puTojY/ozlqy2d/SFNcoLIqTFi42g= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.0/go.mod h1:0jp+ltwkf+SwG2fm/PKo8t4y8pJSgOCO4D8Lz3k0aHQ= +github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.4.3 h1:kT6BcZsmMtNkP/iYMcRG+mIEA/IbeiUimXtGmqF39y0= +github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.4.3/go.mod h1:Z8uGua2k4PPaGOYn66pK02rhMrot3Xk3tpBuUFPomZU= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.3 h1:qcxX0JYlgWH3hpPUnd6U0ikcl6LLA9sLkXE2w1fpMvY= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.3/go.mod h1:cLSNEmI45soc+Ef8K/L+8sEA3A3pYFEYf5B5UI+6bH4= +github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.18.3 h1:ZC7Y/XgKUxwqcdhO5LE8P6oGP1eh6xlQReWNKfhvJno= +github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.18.3/go.mod h1:WqfO7M9l9yUAw0HcHaikwRd/H6gzYdz7vjejCA5e2oY= github.com/aws/aws-sdk-go-v2/service/kms v1.30.1 h1:SBn4I0fJXF9FYOVRSVMWuhvEKoAHDikjGpS3wlmw5DE= github.com/aws/aws-sdk-go-v2/service/kms v1.30.1/go.mod h1:2snWQJQUKsbN66vAawJuOGX7dr37pfOq9hb0tZDGIqQ= -github.com/aws/aws-sdk-go-v2/service/s3 v1.53.1 h1:6cnno47Me9bRykw9AEv9zkXE+5or7jz8TsskTTccbgc= -github.com/aws/aws-sdk-go-v2/service/s3 v1.53.1/go.mod h1:qmdkIIAC+GCLASF7R2whgNrJADz0QZPX+Seiw/i4S3o= +github.com/aws/aws-sdk-go-v2/service/lambda v1.64.1 h1:0njE+T0N80Kl2bPfK85Lnz1+dD/xskJduTqfRyREpvY= +github.com/aws/aws-sdk-go-v2/service/lambda v1.64.1/go.mod h1:hr+VpAzvznKumy8q8TFEJfx3Xx+zfK2gDrrWjBqLLPw= +github.com/aws/aws-sdk-go-v2/service/s3 v1.66.2 h1:p9TNFL8bFUMd+38YIpTAXpoxyz0MxC7FlbFEH4P4E1U= +github.com/aws/aws-sdk-go-v2/service/s3 v1.66.2/go.mod h1:fNjyo0Coen9QTwQLWeV6WO2Nytwiu+cCcWaTdKCAqqE= github.com/aws/aws-sdk-go-v2/service/ssm v1.49.2 h1:JcvYXGYiu7ME17irbW6kvWno2LG5i29Ci0UZyWX0IOs= github.com/aws/aws-sdk-go-v2/service/ssm v1.49.2/go.mod h1:loBAHYxz7JyucJvq4xuW9vunu8iCzjNYfSrQg2QEczA= github.com/aws/aws-sdk-go-v2/service/sso v1.20.5 h1:vN8hEbpRnL7+Hopy9dzmRle1xmDc7o8tmY0klsr175w= @@ -141,8 +145,8 @@ github.com/aws/aws-sdk-go-v2/service/ssooidc v1.23.4 h1:Jux+gDDyi1Lruk+KHF91tK2K github.com/aws/aws-sdk-go-v2/service/ssooidc v1.23.4/go.mod h1:mUYPBhaF2lGiukDEjJX2BLRRKTmoUSitGDUgM4tRxak= github.com/aws/aws-sdk-go-v2/service/sts v1.28.6 h1:cwIxeBttqPN3qkaAjcEcsh8NYr8n2HZPkcKgPAi1phU= github.com/aws/aws-sdk-go-v2/service/sts v1.28.6/go.mod h1:FZf1/nKNEkHdGGJP/cI2MoIMquumuRK6ol3QQJNDxmw= -github.com/aws/smithy-go v1.20.2 h1:tbp628ireGtzcHDDmLT/6ADHidqnwgF57XOXZe6tp4Q= -github.com/aws/smithy-go v1.20.2/go.mod h1:krry+ya/rV9RDcV/Q16kpu6ypI4K2czasz0NC3qS14E= +github.com/aws/smithy-go v1.22.0 h1:uunKnWlcoL3zO7q+gG2Pk53joueEOsnNB28QdMsmiMM= +github.com/aws/smithy-go v1.22.0/go.mod h1:irrKGvNn1InZwb2d7fkIRNucdfwR8R+Ts3wxYa/cJHg= github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k= github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8= github.com/blang/semver v3.5.1+incompatible h1:cQNTCjp13qL8KC3Nbxr/y2Bqb63oX6wdnnjpJbkM4JQ= diff --git a/provider/pkg/client/lambda.go b/provider/pkg/client/lambda.go new file mode 100644 index 0000000000..c92f7928b4 --- /dev/null +++ b/provider/pkg/client/lambda.go @@ -0,0 +1,93 @@ +// Copyright 2016-2024, Pulumi Corporation. + +package client + +import ( + "context" + "fmt" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/lambda" + lambdaTypes "github.com/aws/aws-sdk-go-v2/service/lambda/types" + "github.com/pkg/errors" +) + +const DefaultFunctionActivationTimeout = 5 * time.Minute + +//go:generate mockgen -package client -source lambda.go -destination mock_lambda.go LambdaClient,LambdaApi +type LambdaClient interface { + // InvokeAsync invokes the given Lambda function with the 'Event' invocation type. This means that the function + // will be invoked asynchronously and no response will be returned. If the function is not ready to be invoked yet, + // this method will wait for the function to become active before invoking it. + InvokeAsync(ctx context.Context, functionName string, payload []byte) error +} + +type LambdaApi interface { + Invoke(ctx context.Context, params *lambda.InvokeInput, optFns ...func(*lambda.Options)) (*lambda.InvokeOutput, error) + GetFunction(context.Context, *lambda.GetFunctionInput, ...func(*lambda.Options)) (*lambda.GetFunctionOutput, error) +} + +type lambdaClientImpl struct { + api LambdaApi + functionActivationTimeout time.Duration +} + +func NewLambdaClient(api LambdaApi) LambdaClient { + return &lambdaClientImpl{ + api: api, + functionActivationTimeout: DefaultFunctionActivationTimeout, + } +} + +func (c *lambdaClientImpl) InvokeAsync(ctx context.Context, functionName string, payload []byte) error { + input := lambda.InvokeInput{ + FunctionName: aws.String(functionName), + Payload: payload, + // async invocation + InvocationType: lambdaTypes.InvocationTypeEvent, + } + + // fire off an initial invoke. If the function is not ready, we need to wait for it to become ready. + // this initial invoke will trigger the AWS Lambda service to start the function transition process. + invokeResp, err := c.api.Invoke(ctx, &input) + + if err != nil { + // Lambda functions can be in a state where they are not ready to be invoked yet. + // If we get this error, we need to wait for the function to become active + var notReadyErr *lambdaTypes.ResourceNotReadyException + if errors.As(err, ¬ReadyErr) { + err := c.waitForFunctionActive(ctx, functionName) + if err != nil { + return fmt.Errorf("failed to wait for function to become active: %w", err) + } + invokeResp, err = c.api.Invoke(ctx, &input) + if err != nil { + return err + } + } else { + return err + } + } + + if invokeResp.StatusCode != 202 { + return fmt.Errorf("lambda invocation failed with status code %d", invokeResp.StatusCode) + } + + return nil +} + +// waitForFunctionActive waits for the function to be in the active state. If the function is not ready +// after 5 minutes, it will return an error. +func (c *lambdaClientImpl) waitForFunctionActive(ctx context.Context, functionName string) error { + err := lambda.NewFunctionActiveV2Waiter(c.api, func(o *lambda.FunctionActiveV2WaiterOptions) { + // We already aggressively retry SDK errors with plenty retry attempts and + // throttling exceptions will be taken care of by the SDK + o.MinDelay = time.Second + o.MaxDelay = 5 * time.Second + }).Wait(ctx, &lambda.GetFunctionInput{ + FunctionName: aws.String(functionName), + }, c.functionActivationTimeout) + + return err +} diff --git a/provider/pkg/client/lambda_test.go b/provider/pkg/client/lambda_test.go new file mode 100644 index 0000000000..8811d54c96 --- /dev/null +++ b/provider/pkg/client/lambda_test.go @@ -0,0 +1,150 @@ +package client + +import ( + "context" + "testing" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/lambda" + lambdaTypes "github.com/aws/aws-sdk-go-v2/service/lambda/types" + "github.com/stretchr/testify/assert" + gomock "go.uber.org/mock/gomock" +) + +func lambdaSetup(t *testing.T) (*gomock.Controller, *lambdaClientImpl, *MockLambdaApi) { + ctrl := gomock.NewController(t) + mockLambdaApi := NewMockLambdaApi(ctrl) + client := &lambdaClientImpl{ + api: mockLambdaApi, + functionActivationTimeout: 2 * time.Second, + } + return ctrl, client, mockLambdaApi +} + +func TestInvokeAsync_SuccessfulInvocation(t *testing.T) { + t.Parallel() + ctrl, client, mockLambdaApi := lambdaSetup(t) + defer ctrl.Finish() + + ctx := context.Background() + functionName := "test-function" + payload := []byte(`{"key": "value"}`) + + mockLambdaApi.EXPECT().Invoke(ctx, gomock.Any()).Return(&lambda.InvokeOutput{ + StatusCode: 202, + }, nil) + + err := client.InvokeAsync(ctx, functionName, payload) + assert.NoError(t, err) +} + +func TestInvokeAsync_FunctionNotReadyInitiallyButBecomesReady(t *testing.T) { + t.Parallel() + ctrl, client, mockLambdaApi := lambdaSetup(t) + defer ctrl.Finish() + + ctx := context.Background() + functionName := "test-function" + payload := []byte(`{"key": "value"}`) + + mockLambdaApi.EXPECT().Invoke(ctx, gomock.Any()).Return(nil, &lambdaTypes.ResourceNotReadyException{}) + + mockLambdaApi.EXPECT().GetFunction(gomock.Any(), &lambda.GetFunctionInput{ + FunctionName: aws.String(functionName), + }, gomock.Any()).Return(&lambda.GetFunctionOutput{ + Configuration: &lambdaTypes.FunctionConfiguration{ + State: lambdaTypes.StateInactive, + }, + }, nil) + mockLambdaApi.EXPECT().GetFunction(gomock.Any(), &lambda.GetFunctionInput{ + FunctionName: aws.String(functionName), + }, gomock.Any()).Return(&lambda.GetFunctionOutput{ + Configuration: &lambdaTypes.FunctionConfiguration{ + State: lambdaTypes.StateActive, + }, + }, nil) + + mockLambdaApi.EXPECT().Invoke(ctx, gomock.Any()).Return(&lambda.InvokeOutput{ + StatusCode: 202, + }, nil) + + err := client.InvokeAsync(ctx, functionName, payload) + assert.NoError(t, err) +} + +func TestInvokeAsync_FunctionNotReadyAndFailsToBecomeReady(t *testing.T) { + t.Parallel() + ctrl, client, mockLambdaApi := lambdaSetup(t) + defer ctrl.Finish() + + ctx := context.Background() + functionName := "test-function" + payload := []byte(`{"key": "value"}`) + + mockLambdaApi.EXPECT().Invoke(ctx, gomock.Any()).Return(nil, &lambdaTypes.ResourceNotReadyException{}) + + mockLambdaApi.EXPECT().GetFunction(gomock.Any(), &lambda.GetFunctionInput{ + FunctionName: aws.String(functionName), + }, gomock.Any()).Return(&lambda.GetFunctionOutput{ + Configuration: &lambdaTypes.FunctionConfiguration{ + State: lambdaTypes.StateInactive, + }, + }, nil).AnyTimes() + + err := client.InvokeAsync(ctx, functionName, payload) + assert.Error(t, err) + assert.Contains(t, err.Error(), "failed to wait for function to become active") +} + +func TestInvokeAsync_InvocationFailsWithNon202StatusCode(t *testing.T) { + t.Parallel() + ctrl, client, mockLambdaApi := lambdaSetup(t) + defer ctrl.Finish() + + ctx := context.Background() + functionName := "test-function" + payload := []byte(`{"key": "value"}`) + + mockLambdaApi.EXPECT().Invoke(ctx, gomock.Any()).Return(&lambda.InvokeOutput{ + StatusCode: 400, + }, nil) + + err := client.InvokeAsync(ctx, functionName, payload) + assert.Error(t, err) + assert.Contains(t, err.Error(), "lambda invocation failed with status code 400") +} + +func TestInvokeAsync_FunctionNotReadyInitiallyButBecomesReadyAndThenFails(t *testing.T) { + t.Parallel() + ctrl, client, mockLambdaApi := lambdaSetup(t) + defer ctrl.Finish() + + ctx := context.Background() + functionName := "test-function" + payload := []byte(`{"key": "value"}`) + + mockLambdaApi.EXPECT().Invoke(ctx, gomock.Any()).Return(nil, &lambdaTypes.ResourceNotReadyException{}) + mockLambdaApi.EXPECT().Invoke(ctx, gomock.Any()).Return(&lambda.InvokeOutput{ + StatusCode: 400, + }, nil) + + mockLambdaApi.EXPECT().GetFunction(gomock.Any(), &lambda.GetFunctionInput{ + FunctionName: aws.String(functionName), + }, gomock.Any()).Return(&lambda.GetFunctionOutput{ + Configuration: &lambdaTypes.FunctionConfiguration{ + State: lambdaTypes.StateInactive, + }, + }, nil) + mockLambdaApi.EXPECT().GetFunction(gomock.Any(), &lambda.GetFunctionInput{ + FunctionName: aws.String(functionName), + }, gomock.Any()).Return(&lambda.GetFunctionOutput{ + Configuration: &lambdaTypes.FunctionConfiguration{ + State: lambdaTypes.StateActive, + }, + }, nil) + + err := client.InvokeAsync(ctx, functionName, payload) + assert.Error(t, err) + assert.Contains(t, err.Error(), "lambda invocation failed with status code 400") +} diff --git a/provider/pkg/client/mock_lambda.go b/provider/pkg/client/mock_lambda.go new file mode 100644 index 0000000000..3731257bb4 --- /dev/null +++ b/provider/pkg/client/mock_lambda.go @@ -0,0 +1,120 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: lambda.go +// +// Generated by this command: +// +// mockgen -package client -source lambda.go -destination mock_lambda.go LambdaClient,LambdaApi +// + +// Package client is a generated GoMock package. +package client + +import ( + context "context" + reflect "reflect" + + lambda "github.com/aws/aws-sdk-go-v2/service/lambda" + gomock "go.uber.org/mock/gomock" +) + +// MockLambdaClient is a mock of LambdaClient interface. +type MockLambdaClient struct { + ctrl *gomock.Controller + recorder *MockLambdaClientMockRecorder + isgomock struct{} +} + +// MockLambdaClientMockRecorder is the mock recorder for MockLambdaClient. +type MockLambdaClientMockRecorder struct { + mock *MockLambdaClient +} + +// NewMockLambdaClient creates a new mock instance. +func NewMockLambdaClient(ctrl *gomock.Controller) *MockLambdaClient { + mock := &MockLambdaClient{ctrl: ctrl} + mock.recorder = &MockLambdaClientMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockLambdaClient) EXPECT() *MockLambdaClientMockRecorder { + return m.recorder +} + +// InvokeAsync mocks base method. +func (m *MockLambdaClient) InvokeAsync(ctx context.Context, functionName string, payload []byte) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "InvokeAsync", ctx, functionName, payload) + ret0, _ := ret[0].(error) + return ret0 +} + +// InvokeAsync indicates an expected call of InvokeAsync. +func (mr *MockLambdaClientMockRecorder) InvokeAsync(ctx, functionName, payload any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InvokeAsync", reflect.TypeOf((*MockLambdaClient)(nil).InvokeAsync), ctx, functionName, payload) +} + +// MockLambdaApi is a mock of LambdaApi interface. +type MockLambdaApi struct { + ctrl *gomock.Controller + recorder *MockLambdaApiMockRecorder + isgomock struct{} +} + +// MockLambdaApiMockRecorder is the mock recorder for MockLambdaApi. +type MockLambdaApiMockRecorder struct { + mock *MockLambdaApi +} + +// NewMockLambdaApi creates a new mock instance. +func NewMockLambdaApi(ctrl *gomock.Controller) *MockLambdaApi { + mock := &MockLambdaApi{ctrl: ctrl} + mock.recorder = &MockLambdaApiMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockLambdaApi) EXPECT() *MockLambdaApiMockRecorder { + return m.recorder +} + +// GetFunction mocks base method. +func (m *MockLambdaApi) GetFunction(arg0 context.Context, arg1 *lambda.GetFunctionInput, arg2 ...func(*lambda.Options)) (*lambda.GetFunctionOutput, error) { + m.ctrl.T.Helper() + varargs := []any{arg0, arg1} + for _, a := range arg2 { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "GetFunction", varargs...) + ret0, _ := ret[0].(*lambda.GetFunctionOutput) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetFunction indicates an expected call of GetFunction. +func (mr *MockLambdaApiMockRecorder) GetFunction(arg0, arg1 any, arg2 ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{arg0, arg1}, arg2...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetFunction", reflect.TypeOf((*MockLambdaApi)(nil).GetFunction), varargs...) +} + +// Invoke mocks base method. +func (m *MockLambdaApi) Invoke(ctx context.Context, params *lambda.InvokeInput, optFns ...func(*lambda.Options)) (*lambda.InvokeOutput, error) { + m.ctrl.T.Helper() + varargs := []any{ctx, params} + for _, a := range optFns { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "Invoke", varargs...) + ret0, _ := ret[0].(*lambda.InvokeOutput) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Invoke indicates an expected call of Invoke. +func (mr *MockLambdaApiMockRecorder) Invoke(ctx, params any, optFns ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{ctx, params}, optFns...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Invoke", reflect.TypeOf((*MockLambdaApi)(nil).Invoke), varargs...) +} diff --git a/provider/pkg/client/mock_s3.go b/provider/pkg/client/mock_s3.go new file mode 100644 index 0000000000..71676d0493 --- /dev/null +++ b/provider/pkg/client/mock_s3.go @@ -0,0 +1,183 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: s3.go +// +// Generated by this command: +// +// mockgen -package client -source s3.go -destination mock_s3.go S3Client,S3Api,S3PresignApi +// + +// Package client is a generated GoMock package. +package client + +import ( + context "context" + io "io" + reflect "reflect" + time "time" + + v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4" + s3 "github.com/aws/aws-sdk-go-v2/service/s3" + gomock "go.uber.org/mock/gomock" +) + +// MockS3Client is a mock of S3Client interface. +type MockS3Client struct { + ctrl *gomock.Controller + recorder *MockS3ClientMockRecorder + isgomock struct{} +} + +// MockS3ClientMockRecorder is the mock recorder for MockS3Client. +type MockS3ClientMockRecorder struct { + mock *MockS3Client +} + +// NewMockS3Client creates a new mock instance. +func NewMockS3Client(ctrl *gomock.Controller) *MockS3Client { + mock := &MockS3Client{ctrl: ctrl} + mock.recorder = &MockS3ClientMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockS3Client) EXPECT() *MockS3ClientMockRecorder { + return m.recorder +} + +// PresignPutObject mocks base method. +func (m *MockS3Client) PresignPutObject(ctx context.Context, bucket, key string, expiration time.Duration) (string, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "PresignPutObject", ctx, bucket, key, expiration) + ret0, _ := ret[0].(string) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// PresignPutObject indicates an expected call of PresignPutObject. +func (mr *MockS3ClientMockRecorder) PresignPutObject(ctx, bucket, key, expiration any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PresignPutObject", reflect.TypeOf((*MockS3Client)(nil).PresignPutObject), ctx, bucket, key, expiration) +} + +// WaitForObject mocks base method. +func (m *MockS3Client) WaitForObject(ctx context.Context, bucket, key string, timeout time.Duration) (io.ReadCloser, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "WaitForObject", ctx, bucket, key, timeout) + ret0, _ := ret[0].(io.ReadCloser) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// WaitForObject indicates an expected call of WaitForObject. +func (mr *MockS3ClientMockRecorder) WaitForObject(ctx, bucket, key, timeout any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WaitForObject", reflect.TypeOf((*MockS3Client)(nil).WaitForObject), ctx, bucket, key, timeout) +} + +// MockS3Api is a mock of S3Api interface. +type MockS3Api struct { + ctrl *gomock.Controller + recorder *MockS3ApiMockRecorder + isgomock struct{} +} + +// MockS3ApiMockRecorder is the mock recorder for MockS3Api. +type MockS3ApiMockRecorder struct { + mock *MockS3Api +} + +// NewMockS3Api creates a new mock instance. +func NewMockS3Api(ctrl *gomock.Controller) *MockS3Api { + mock := &MockS3Api{ctrl: ctrl} + mock.recorder = &MockS3ApiMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockS3Api) EXPECT() *MockS3ApiMockRecorder { + return m.recorder +} + +// GetObject mocks base method. +func (m *MockS3Api) GetObject(ctx context.Context, params *s3.GetObjectInput, optFns ...func(*s3.Options)) (*s3.GetObjectOutput, error) { + m.ctrl.T.Helper() + varargs := []any{ctx, params} + for _, a := range optFns { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "GetObject", varargs...) + ret0, _ := ret[0].(*s3.GetObjectOutput) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetObject indicates an expected call of GetObject. +func (mr *MockS3ApiMockRecorder) GetObject(ctx, params any, optFns ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{ctx, params}, optFns...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetObject", reflect.TypeOf((*MockS3Api)(nil).GetObject), varargs...) +} + +// HeadObject mocks base method. +func (m *MockS3Api) HeadObject(arg0 context.Context, arg1 *s3.HeadObjectInput, arg2 ...func(*s3.Options)) (*s3.HeadObjectOutput, error) { + m.ctrl.T.Helper() + varargs := []any{arg0, arg1} + for _, a := range arg2 { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "HeadObject", varargs...) + ret0, _ := ret[0].(*s3.HeadObjectOutput) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// HeadObject indicates an expected call of HeadObject. +func (mr *MockS3ApiMockRecorder) HeadObject(arg0, arg1 any, arg2 ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{arg0, arg1}, arg2...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HeadObject", reflect.TypeOf((*MockS3Api)(nil).HeadObject), varargs...) +} + +// MockS3PresignApi is a mock of S3PresignApi interface. +type MockS3PresignApi struct { + ctrl *gomock.Controller + recorder *MockS3PresignApiMockRecorder + isgomock struct{} +} + +// MockS3PresignApiMockRecorder is the mock recorder for MockS3PresignApi. +type MockS3PresignApiMockRecorder struct { + mock *MockS3PresignApi +} + +// NewMockS3PresignApi creates a new mock instance. +func NewMockS3PresignApi(ctrl *gomock.Controller) *MockS3PresignApi { + mock := &MockS3PresignApi{ctrl: ctrl} + mock.recorder = &MockS3PresignApiMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockS3PresignApi) EXPECT() *MockS3PresignApiMockRecorder { + return m.recorder +} + +// PresignPutObject mocks base method. +func (m *MockS3PresignApi) PresignPutObject(ctx context.Context, params *s3.PutObjectInput, optFns ...func(*s3.PresignOptions)) (*v4.PresignedHTTPRequest, error) { + m.ctrl.T.Helper() + varargs := []any{ctx, params} + for _, a := range optFns { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "PresignPutObject", varargs...) + ret0, _ := ret[0].(*v4.PresignedHTTPRequest) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// PresignPutObject indicates an expected call of PresignPutObject. +func (mr *MockS3PresignApiMockRecorder) PresignPutObject(ctx, params any, optFns ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{ctx, params}, optFns...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PresignPutObject", reflect.TypeOf((*MockS3PresignApi)(nil).PresignPutObject), varargs...) +} diff --git a/provider/pkg/client/s3.go b/provider/pkg/client/s3.go new file mode 100644 index 0000000000..68c0235e0f --- /dev/null +++ b/provider/pkg/client/s3.go @@ -0,0 +1,96 @@ +// Copyright 2016-2024, Pulumi Corporation. + +package client + +import ( + "context" + "fmt" + "io" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + signerV4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4" + "github.com/aws/aws-sdk-go-v2/service/s3" + s3Types "github.com/aws/aws-sdk-go-v2/service/s3/types" + "github.com/pkg/errors" +) + +//go:generate mockgen -package client -source s3.go -destination mock_s3.go S3Client,S3Api,S3PresignApi +type S3Client interface { + // PresignPutObject generates a pre-signed URL for uploading an object to an S3 bucket. + // The URL will be valid for the specified expiration duration. + PresignPutObject(ctx context.Context, bucket, key string, expiration time.Duration) (string, error) + + // WaitForObject waits for an object to exist in an S3 bucket and returns a reader for the object. + // The function will block until the object exists or the timeout is reached. + // If the object does not exist after the timeout, an error will be returned. + // The caller is responsible for closing the reader when done. + WaitForObject(ctx context.Context, bucket string, key string, timeout time.Duration) (io.ReadCloser, error) +} + +type S3Api interface { + HeadObject(context.Context, *s3.HeadObjectInput, ...func(*s3.Options)) (*s3.HeadObjectOutput, error) + GetObject(ctx context.Context, params *s3.GetObjectInput, optFns ...func(*s3.Options)) (*s3.GetObjectOutput, error) +} + +type S3PresignApi interface { + PresignPutObject(ctx context.Context, params *s3.PutObjectInput, optFns ...func(*s3.PresignOptions)) (*signerV4.PresignedHTTPRequest, error) +} + +type s3ClientImpl struct { + api S3Api + presignApi S3PresignApi +} + +func NewS3Client(api S3Api, presignApi S3PresignApi) S3Client { + return &s3ClientImpl{ + api: api, + presignApi: presignApi, + } +} + +func (c *s3ClientImpl) PresignPutObject(ctx context.Context, bucket string, key string, expiration time.Duration) (string, error) { + request, err := c.presignApi.PresignPutObject(ctx, &s3.PutObjectInput{ + Bucket: aws.String(bucket), + Key: aws.String(key), + }, func(options *s3.PresignOptions) { + options.Expires = expiration + }) + if err != nil { + return "", err + } + return request.URL, nil +} + +func (c *s3ClientImpl) WaitForObject(ctx context.Context, bucket string, key string, timeout time.Duration) (io.ReadCloser, error) { + err := s3.NewObjectExistsWaiter(c.api, func(o *s3.ObjectExistsWaiterOptions) { + // We already aggressively retry SDK errors with plenty retry attempts and + // throttling exceptions will be taken care of by the SDK + o.MinDelay = time.Second + o.MaxDelay = 5 * time.Second + }).Wait(ctx, &s3.HeadObjectInput{ + Bucket: aws.String(bucket), + Key: aws.String(key), + }, timeout) + + if err != nil { + return nil, errors.Wrapf(err, "failed waiting for object %q in bucket %q", key, bucket) + } + + getResponse, err := c.api.GetObject(ctx, &s3.GetObjectInput{ + Bucket: aws.String(bucket), + Key: aws.String(key), + }) + + if err != nil { + var noKey *s3Types.NoSuchKey + if errors.As(err, &noKey) { + // this could happen if the response object got deleted before we can read it + return nil, fmt.Errorf("object %q in bucket %q was deleted before it was retrieved", key, bucket) + } else { + return nil, errors.Wrapf(err, "failed to get object %q from bucket %q", key, bucket) + } + } + + return getResponse.Body, nil +} diff --git a/provider/pkg/client/s3_test.go b/provider/pkg/client/s3_test.go new file mode 100644 index 0000000000..077dcb741e --- /dev/null +++ b/provider/pkg/client/s3_test.go @@ -0,0 +1,188 @@ +package client + +import ( + "context" + "errors" + io "io" + "strings" + "testing" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4" + "github.com/aws/aws-sdk-go-v2/service/s3" + s3Types "github.com/aws/aws-sdk-go-v2/service/s3/types" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + gomock "go.uber.org/mock/gomock" +) + +func TestPresignPutObject(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockPresignApi := NewMockS3PresignApi(ctrl) + client := &s3ClientImpl{ + presignApi: mockPresignApi, + } + + ctx := context.Background() + bucket := "test-bucket" + key := "test-key" + expiration := 15 * time.Minute + expectedURL := "https://example.com/presigned-url" + + mockPresignApi.EXPECT().PresignPutObject(ctx, &s3.PutObjectInput{ + Bucket: aws.String(bucket), + Key: aws.String(key), + }, gomock.Any()).Return(&v4.PresignedHTTPRequest{ + URL: expectedURL, + }, nil) + + url, err := client.PresignPutObject(ctx, bucket, key, expiration) + assert.NoError(t, err) + assert.Equal(t, expectedURL, url) +} + +func TestPresignPutObject_Error(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockPresignApi := NewMockS3PresignApi(ctrl) + client := &s3ClientImpl{ + presignApi: mockPresignApi, + } + + ctx := context.Background() + bucket := "test-bucket" + key := "test-key" + expiration := 15 * time.Minute + expectedError := errors.New("presign error") + + mockPresignApi.EXPECT().PresignPutObject(ctx, &s3.PutObjectInput{ + Bucket: aws.String(bucket), + Key: aws.String(key), + }, gomock.Any()).Return(nil, expectedError) + + url, err := client.PresignPutObject(ctx, bucket, key, expiration) + assert.Error(t, err) + assert.Equal(t, expectedError, err) + assert.Empty(t, url) +} + +func s3Setup(t *testing.T) (*gomock.Controller, *s3ClientImpl, context.Context) { + ctrl := gomock.NewController(t) + mockApi := NewMockS3Api(ctrl) + client := &s3ClientImpl{ + api: mockApi, + } + ctx := context.Background() + return ctrl, client, ctx +} + +func TestWaitForObject_Success(t *testing.T) { + t.Parallel() + ctrl, client, ctx := s3Setup(t) + defer ctrl.Finish() + + bucket := "test-bucket" + key := "test-key" + timeout := 2 * time.Second + expectedBody := "test content" + + mockApi := client.api.(*MockS3Api) + mockApi.EXPECT().HeadObject(gomock.Any(), &s3.HeadObjectInput{ + Bucket: aws.String(bucket), + Key: aws.String(key), + }, gomock.Any()).Return(&s3.HeadObjectOutput{}, nil) + + mockApi.EXPECT().GetObject(ctx, &s3.GetObjectInput{ + Bucket: aws.String(bucket), + Key: aws.String(key), + }, gomock.Any()).Return(&s3.GetObjectOutput{ + Body: io.NopCloser(strings.NewReader(expectedBody)), + }, nil) + + body, err := client.WaitForObject(ctx, bucket, key, timeout) + require.NoError(t, err) + + defer body.Close() + content, err := io.ReadAll(body) + require.NoError(t, err) + assert.Equal(t, expectedBody, string(content)) +} + +func TestWaitForObject_Timeout(t *testing.T) { + t.Parallel() + ctrl, client, ctx := s3Setup(t) + defer ctrl.Finish() + + bucket := "test-bucket" + key := "test-key" + timeout := 2 * time.Second + + mockApi := client.api.(*MockS3Api) + mockApi.EXPECT().HeadObject(gomock.Any(), &s3.HeadObjectInput{ + Bucket: aws.String(bucket), + Key: aws.String(key), + }, gomock.Any()).Return(nil, &s3Types.NotFound{}).AnyTimes() + + body, err := client.WaitForObject(ctx, bucket, key, timeout) + assert.Error(t, err) + assert.Contains(t, err.Error(), "failed waiting for object") + assert.Nil(t, body) +} + +func TestWaitForObject_Deleted(t *testing.T) { + t.Parallel() + ctrl, client, ctx := s3Setup(t) + defer ctrl.Finish() + + bucket := "test-bucket" + key := "test-key" + timeout := 2 * time.Second + expectedError := &s3Types.NoSuchKey{} + + mockApi := client.api.(*MockS3Api) + mockApi.EXPECT().HeadObject(gomock.Any(), &s3.HeadObjectInput{ + Bucket: aws.String(bucket), + Key: aws.String(key), + }, gomock.Any()).Return(&s3.HeadObjectOutput{}, nil) + + mockApi.EXPECT().GetObject(ctx, &s3.GetObjectInput{ + Bucket: aws.String(bucket), + Key: aws.String(key), + }, gomock.Any()).Return(nil, expectedError) + + body, err := client.WaitForObject(ctx, bucket, key, timeout) + assert.Error(t, err) + assert.Contains(t, err.Error(), "was deleted before it was retrieved") + assert.Nil(t, body) +} + +func TestWaitForObject_GetObjectError(t *testing.T) { + t.Parallel() + ctrl, client, ctx := s3Setup(t) + defer ctrl.Finish() + + bucket := "test-bucket" + key := "test-key" + timeout := 2 * time.Second + expectedError := errors.New("get object error") + + mockApi := client.api.(*MockS3Api) + mockApi.EXPECT().HeadObject(gomock.Any(), &s3.HeadObjectInput{ + Bucket: aws.String(bucket), + Key: aws.String(key), + }, gomock.Any()).Return(&s3.HeadObjectOutput{}, nil) + + mockApi.EXPECT().GetObject(ctx, &s3.GetObjectInput{ + Bucket: aws.String(bucket), + Key: aws.String(key), + }, gomock.Any()).Return(nil, expectedError) + + body, err := client.WaitForObject(ctx, bucket, key, timeout) + assert.Error(t, err) + assert.Contains(t, err.Error(), "failed to get object") + assert.Nil(t, body) +}