Skip to content

Commit 78df77d

Browse files
bbernaysbbernays
authored
feat: Test MeterUsage API call on initial setup of client (#1906)
 This allows customers to validate that the credentials they are using are valid for reporting usage
1 parent bfce6fe commit 78df77d

3 files changed

Lines changed: 36 additions & 5 deletions

File tree

‎go.mod‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ require (
1010
github.com/aws/aws-sdk-go-v2/config v1.28.0
1111
github.com/aws/aws-sdk-go-v2/service/licensemanager v1.29.2
1212
github.com/aws/aws-sdk-go-v2/service/marketplacemetering v1.25.2
13+
github.com/aws/smithy-go v1.22.0
1314
github.com/bradleyjkemp/cupaloy/v2 v2.8.0
1415
github.com/cloudquery/cloudquery-api-go v1.13.1
1516
github.com/cloudquery/plugin-pb-go v1.25.0
@@ -60,7 +61,6 @@ require (
6061
github.com/aws/aws-sdk-go-v2/service/sso v1.24.2 // indirect
6162
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.28.2 // indirect
6263
github.com/aws/aws-sdk-go-v2/service/sts v1.32.2 // indirect
63-
github.com/aws/smithy-go v1.22.0 // indirect
6464
github.com/bahlo/generic-list-go v0.2.0 // indirect
6565
github.com/buger/jsonparser v1.1.1 // indirect
6666
github.com/cenkalti/backoff/v4 v4.3.0 // indirect

‎premium/usage.go‎

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import (
1515
awsConfig "github.com/aws/aws-sdk-go-v2/config"
1616
"github.com/aws/aws-sdk-go-v2/service/marketplacemetering"
1717
"github.com/aws/aws-sdk-go-v2/service/marketplacemetering/types"
18+
"github.com/aws/smithy-go"
1819
cqapi "github.com/cloudquery/cloudquery-api-go"
1920
"github.com/cloudquery/cloudquery-api-go/auth"
2021
"github.com/cloudquery/cloudquery-api-go/config"
@@ -274,7 +275,8 @@ func NewUsageClient(meta plugin.Meta, ops ...UsageClientOptions) (UsageClient, e
274275
}
275276

276277
func (u *BatchUpdater) setupAWSMarketplace() error {
277-
cfg, err := awsConfig.LoadDefaultConfig(context.TODO())
278+
ctx := context.TODO()
279+
cfg, err := awsConfig.LoadDefaultConfig(ctx)
278280
if err != nil {
279281
return fmt.Errorf("failed to load AWS config: %w", err)
280282
}
@@ -288,7 +290,19 @@ func (u *BatchUpdater) setupAWSMarketplace() error {
288290

289291
u.minTimeBetweenFlushes = 1 * time.Minute
290292
u.backgroundUpdater()
291-
return nil
293+
294+
_, err = u.awsMarketplaceClient.MeterUsage(ctx, &marketplacemetering.MeterUsageInput{
295+
ProductCode: aws.String(awsMarketplaceProductCode()),
296+
Timestamp: aws.Time(time.Now()),
297+
UsageDimension: aws.String("rows"),
298+
UsageQuantity: aws.Int32(int32(0)),
299+
DryRun: aws.Bool(true),
300+
})
301+
var apiErr smithy.APIError
302+
if errors.As(err, &apiErr) && apiErr.ErrorCode() == "DryRunOperation" {
303+
return nil
304+
}
305+
return fmt.Errorf("failed dry run invocation with error: %w", err)
292306
}
293307

294308
func isAWSMarketplace() bool {

‎premium/usage_test.go‎

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import (
1515
"github.com/aws/aws-sdk-go-v2/aws"
1616
"github.com/aws/aws-sdk-go-v2/service/marketplacemetering"
1717
"github.com/aws/aws-sdk-go-v2/service/marketplacemetering/types"
18+
"github.com/aws/smithy-go"
1819
cqapi "github.com/cloudquery/cloudquery-api-go"
1920
"github.com/cloudquery/cloudquery-api-go/auth"
2021
"github.com/cloudquery/cloudquery-api-go/config"
@@ -366,6 +367,15 @@ func TestUsageService_AWSMarketplaceDone(t *testing.T) {
366367
m := mocks.NewMockAWSMarketplaceClientInterface(ctrl)
367368

368369
out := marketplacemetering.MeterUsageOutput{}
370+
inTest := meteringInput{
371+
marketplacemetering.MeterUsageInput{
372+
ProductCode: aws.String("2a8bdkarwqrp0tmo4errl65s7"),
373+
UsageDimension: aws.String("rows"),
374+
UsageQuantity: aws.Int32(int32(0)),
375+
DryRun: aws.Bool(true)},
376+
}
377+
errTest := smithy.GenericAPIError{Code: "DryRunOperation", Message: "No errors detected in dry run"}
378+
369379
in := meteringInput{
370380
MeterUsageInput: marketplacemetering.MeterUsageInput{
371381
ProductCode: aws.String("2a8bdkarwqrp0tmo4errl65s7"),
@@ -391,7 +401,12 @@ func TestUsageService_AWSMarketplaceDone(t *testing.T) {
391401
},
392402
}
393403
assert.NoError(t, faker.FakeObject(&out))
394-
m.EXPECT().MeterUsage(gomock.Any(), in).Return(&out, nil)
404+
405+
gomock.InOrder(
406+
m.EXPECT().MeterUsage(gomock.Any(), inTest).Return(&out, &errTest),
407+
m.EXPECT().MeterUsage(gomock.Any(), in).Return(&out, nil),
408+
)
409+
395410
t.Setenv("CQ_AWS_MARKETPLACE_CONTAINER", "true")
396411
usageClient := newClient(t, nil, WithBatchLimit(50), WithAWSMarketplaceClient(m))
397412

@@ -909,7 +924,9 @@ func (mi meteringInput) Matches(x any) bool {
909924
if aws.ToInt32(testInput.UsageQuantity) != aws.ToInt32(mi.UsageQuantity) {
910925
return false
911926
}
912-
927+
if aws.ToBool(testInput.DryRun) != aws.ToBool(mi.DryRun) {
928+
return false
929+
}
913930
return true
914931
}
915932

0 commit comments

Comments
 (0)