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
cloudquery · Oct 31, 2024 · 78df77d · 78df77d
1 parent bfce6fe
commit 78df77d
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 5 deletions.
diff --git a/go.mod b/go.mod
@@ -10,6 +10,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/config v1.28.0
 	github.com/aws/aws-sdk-go-v2/service/licensemanager v1.29.2
 	github.com/aws/aws-sdk-go-v2/service/marketplacemetering v1.25.2
+	github.com/aws/smithy-go v1.22.0
 	github.com/bradleyjkemp/cupaloy/v2 v2.8.0
 	github.com/cloudquery/cloudquery-api-go v1.13.1
 	github.com/cloudquery/plugin-pb-go v1.25.0
@@ -60,7 +61,6 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/sso v1.24.2 // indirect
 	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.28.2 // indirect
 	github.com/aws/aws-sdk-go-v2/service/sts v1.32.2 // indirect
-	github.com/aws/smithy-go v1.22.0 // indirect
 	github.com/bahlo/generic-list-go v0.2.0 // indirect
 	github.com/buger/jsonparser v1.1.1 // indirect
 	github.com/cenkalti/backoff/v4 v4.3.0 // indirect

diff --git a/premium/usage.go b/premium/usage.go
@@ -15,6 +15,7 @@ import (
 	awsConfig "github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/service/marketplacemetering"
 	"github.com/aws/aws-sdk-go-v2/service/marketplacemetering/types"
+	"github.com/aws/smithy-go"
 	cqapi "github.com/cloudquery/cloudquery-api-go"
 	"github.com/cloudquery/cloudquery-api-go/auth"
 	"github.com/cloudquery/cloudquery-api-go/config"
@@ -274,7 +275,8 @@ func NewUsageClient(meta plugin.Meta, ops ...UsageClientOptions) (UsageClient, e
 }
 
 func (u *BatchUpdater) setupAWSMarketplace() error {
-	cfg, err := awsConfig.LoadDefaultConfig(context.TODO())
+	ctx := context.TODO()
+	cfg, err := awsConfig.LoadDefaultConfig(ctx)
 	if err != nil {
 		return fmt.Errorf("failed to load AWS config: %w", err)
 	}
@@ -288,7 +290,19 @@ func (u *BatchUpdater) setupAWSMarketplace() error {
 
 	u.minTimeBetweenFlushes = 1 * time.Minute
 	u.backgroundUpdater()
-	return nil
+
+	_, err = u.awsMarketplaceClient.MeterUsage(ctx, &marketplacemetering.MeterUsageInput{
+		ProductCode:    aws.String(awsMarketplaceProductCode()),
+		Timestamp:      aws.Time(time.Now()),
+		UsageDimension: aws.String("rows"),
+		UsageQuantity:  aws.Int32(int32(0)),
+		DryRun:         aws.Bool(true),
+	})
+	var apiErr smithy.APIError
+	if errors.As(err, &apiErr) && apiErr.ErrorCode() == "DryRunOperation" {
+		return nil
+	}
+	return fmt.Errorf("failed dry run invocation with error: %w", err)
 }
 
 func isAWSMarketplace() bool {

diff --git a/premium/usage_test.go b/premium/usage_test.go
@@ -15,6 +15,7 @@ import (
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/marketplacemetering"
 	"github.com/aws/aws-sdk-go-v2/service/marketplacemetering/types"
+	"github.com/aws/smithy-go"
 	cqapi "github.com/cloudquery/cloudquery-api-go"
 	"github.com/cloudquery/cloudquery-api-go/auth"
 	"github.com/cloudquery/cloudquery-api-go/config"
@@ -366,6 +367,15 @@ func TestUsageService_AWSMarketplaceDone(t *testing.T) {
 	m := mocks.NewMockAWSMarketplaceClientInterface(ctrl)
 
 	out := marketplacemetering.MeterUsageOutput{}
+	inTest := meteringInput{
+		marketplacemetering.MeterUsageInput{
+			ProductCode:    aws.String("2a8bdkarwqrp0tmo4errl65s7"),
+			UsageDimension: aws.String("rows"),
+			UsageQuantity:  aws.Int32(int32(0)),
+			DryRun:         aws.Bool(true)},
+	}
+	errTest := smithy.GenericAPIError{Code: "DryRunOperation", Message: "No errors detected in dry run"}
+
 	in := meteringInput{
 		MeterUsageInput: marketplacemetering.MeterUsageInput{
 			ProductCode:    aws.String("2a8bdkarwqrp0tmo4errl65s7"),
@@ -391,7 +401,12 @@ func TestUsageService_AWSMarketplaceDone(t *testing.T) {
 		},
 	}
 	assert.NoError(t, faker.FakeObject(&out))
-	m.EXPECT().MeterUsage(gomock.Any(), in).Return(&out, nil)
+
+	gomock.InOrder(
+		m.EXPECT().MeterUsage(gomock.Any(), inTest).Return(&out, &errTest),
+		m.EXPECT().MeterUsage(gomock.Any(), in).Return(&out, nil),
+	)
+
 	t.Setenv("CQ_AWS_MARKETPLACE_CONTAINER", "true")
 	usageClient := newClient(t, nil, WithBatchLimit(50), WithAWSMarketplaceClient(m))
 
@@ -909,7 +924,9 @@ func (mi meteringInput) Matches(x any) bool {
 	if aws.ToInt32(testInput.UsageQuantity) != aws.ToInt32(mi.UsageQuantity) {
 		return false
 	}
-
+	if aws.ToBool(testInput.DryRun) != aws.ToBool(mi.DryRun) {
+		return false
+	}
 	return true
 }