From ad470ec285ec7d482ad66ab7e2a54bc960350091 Mon Sep 17 00:00:00 2001 From: pokom Date: Tue, 5 Dec 2023 14:50:06 -0500 Subject: [PATCH] Refactor billing client --- mocks/pkg/google/gcs/CloudCatalogClient.go | 169 +++++++++++++++++++++ pkg/google/gcp.go | 13 +- pkg/google/gcs/bucket_test.go | 4 +- pkg/google/gcs/gcs.go | 58 +++---- pkg/google/gcs/gcs_test.go | 9 +- 5 files changed, 212 insertions(+), 41 deletions(-) create mode 100644 mocks/pkg/google/gcs/CloudCatalogClient.go diff --git a/mocks/pkg/google/gcs/CloudCatalogClient.go b/mocks/pkg/google/gcs/CloudCatalogClient.go new file mode 100644 index 00000000..271946cf --- /dev/null +++ b/mocks/pkg/google/gcs/CloudCatalogClient.go @@ -0,0 +1,169 @@ +// Code generated by mockery v2.38.0. DO NOT EDIT. + +package gcs + +import ( + billing "cloud.google.com/go/billing/apiv1" + billingpb "cloud.google.com/go/billing/apiv1/billingpb" + + context "context" + + gax "github.com/googleapis/gax-go/v2" + + mock "github.com/stretchr/testify/mock" +) + +// CloudCatalogClient is an autogenerated mock type for the CloudCatalogClient type +type CloudCatalogClient struct { + mock.Mock +} + +type CloudCatalogClient_Expecter struct { + mock *mock.Mock +} + +func (_m *CloudCatalogClient) EXPECT() *CloudCatalogClient_Expecter { + return &CloudCatalogClient_Expecter{mock: &_m.Mock} +} + +// ListServices provides a mock function with given fields: ctx, req, opts +func (_m *CloudCatalogClient) ListServices(ctx context.Context, req *billingpb.ListServicesRequest, opts ...gax.CallOption) *billing.ServiceIterator { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, req) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + if len(ret) == 0 { + panic("no return value specified for ListServices") + } + + var r0 *billing.ServiceIterator + if rf, ok := ret.Get(0).(func(context.Context, *billingpb.ListServicesRequest, ...gax.CallOption) *billing.ServiceIterator); ok { + r0 = rf(ctx, req, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*billing.ServiceIterator) + } + } + + return r0 +} + +// CloudCatalogClient_ListServices_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ListServices' +type CloudCatalogClient_ListServices_Call struct { + *mock.Call +} + +// ListServices is a helper method to define mock.On call +// - ctx context.Context +// - req *billingpb.ListServicesRequest +// - opts ...gax.CallOption +func (_e *CloudCatalogClient_Expecter) ListServices(ctx interface{}, req interface{}, opts ...interface{}) *CloudCatalogClient_ListServices_Call { + return &CloudCatalogClient_ListServices_Call{Call: _e.mock.On("ListServices", + append([]interface{}{ctx, req}, opts...)...)} +} + +func (_c *CloudCatalogClient_ListServices_Call) Run(run func(ctx context.Context, req *billingpb.ListServicesRequest, opts ...gax.CallOption)) *CloudCatalogClient_ListServices_Call { + _c.Call.Run(func(args mock.Arguments) { + variadicArgs := make([]gax.CallOption, len(args)-2) + for i, a := range args[2:] { + if a != nil { + variadicArgs[i] = a.(gax.CallOption) + } + } + run(args[0].(context.Context), args[1].(*billingpb.ListServicesRequest), variadicArgs...) + }) + return _c +} + +func (_c *CloudCatalogClient_ListServices_Call) Return(_a0 *billing.ServiceIterator) *CloudCatalogClient_ListServices_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *CloudCatalogClient_ListServices_Call) RunAndReturn(run func(context.Context, *billingpb.ListServicesRequest, ...gax.CallOption) *billing.ServiceIterator) *CloudCatalogClient_ListServices_Call { + _c.Call.Return(run) + return _c +} + +// ListSkus provides a mock function with given fields: ctx, req, opts +func (_m *CloudCatalogClient) ListSkus(ctx context.Context, req *billingpb.ListSkusRequest, opts ...gax.CallOption) *billing.SkuIterator { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, req) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + if len(ret) == 0 { + panic("no return value specified for ListSkus") + } + + var r0 *billing.SkuIterator + if rf, ok := ret.Get(0).(func(context.Context, *billingpb.ListSkusRequest, ...gax.CallOption) *billing.SkuIterator); ok { + r0 = rf(ctx, req, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*billing.SkuIterator) + } + } + + return r0 +} + +// CloudCatalogClient_ListSkus_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ListSkus' +type CloudCatalogClient_ListSkus_Call struct { + *mock.Call +} + +// ListSkus is a helper method to define mock.On call +// - ctx context.Context +// - req *billingpb.ListSkusRequest +// - opts ...gax.CallOption +func (_e *CloudCatalogClient_Expecter) ListSkus(ctx interface{}, req interface{}, opts ...interface{}) *CloudCatalogClient_ListSkus_Call { + return &CloudCatalogClient_ListSkus_Call{Call: _e.mock.On("ListSkus", + append([]interface{}{ctx, req}, opts...)...)} +} + +func (_c *CloudCatalogClient_ListSkus_Call) Run(run func(ctx context.Context, req *billingpb.ListSkusRequest, opts ...gax.CallOption)) *CloudCatalogClient_ListSkus_Call { + _c.Call.Run(func(args mock.Arguments) { + variadicArgs := make([]gax.CallOption, len(args)-2) + for i, a := range args[2:] { + if a != nil { + variadicArgs[i] = a.(gax.CallOption) + } + } + run(args[0].(context.Context), args[1].(*billingpb.ListSkusRequest), variadicArgs...) + }) + return _c +} + +func (_c *CloudCatalogClient_ListSkus_Call) Return(_a0 *billing.SkuIterator) *CloudCatalogClient_ListSkus_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *CloudCatalogClient_ListSkus_Call) RunAndReturn(run func(context.Context, *billingpb.ListSkusRequest, ...gax.CallOption) *billing.SkuIterator) *CloudCatalogClient_ListSkus_Call { + _c.Call.Return(run) + return _c +} + +// NewCloudCatalogClient creates a new instance of CloudCatalogClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewCloudCatalogClient(t interface { + mock.TestingT + Cleanup(func()) +}) *CloudCatalogClient { + mock := &CloudCatalogClient{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/pkg/google/gcp.go b/pkg/google/gcp.go index 33e9f804..a793ba5c 100644 --- a/pkg/google/gcp.go +++ b/pkg/google/gcp.go @@ -44,9 +44,9 @@ func New(config *Config) (*GCP, error) { return nil, fmt.Errorf("error creating compute computeService: %w", err) } - billingService, err := billingv1.NewCloudCatalogClient(ctx) + cloudCatalogClient, err := billingv1.NewCloudCatalogClient(ctx) if err != nil { - return nil, fmt.Errorf("error creating billing billingService: %w", err) + return nil, fmt.Errorf("error creating cloudCatalogClient: %w", err) } regionsClient, err := computeapiv1.NewRegionsRESTClient(ctx) @@ -65,12 +65,17 @@ func New(config *Config) (*GCP, error) { var collector provider.Collector switch strings.ToUpper(service) { case "GCS": + serviceName, err := gcs.GetServiceNameByReadableName(ctx, cloudCatalogClient, "Cloud Storage") + if err != nil { + return nil, fmt.Errorf("could not get service name for GCS: %v", err) + } collector, err = gcs.New(&gcs.Config{ ProjectId: config.ProjectId, Projects: config.Projects, ScrapeInterval: config.ScrapeInterval, DefaultDiscount: config.DefaultDiscount, - }, billingService, regionsClient, storageClient) + ServiceName: serviceName, + }, cloudCatalogClient, regionsClient, storageClient) if err != nil { log.Printf("Error creating GCS collector: %s", err) continue @@ -78,7 +83,7 @@ func New(config *Config) (*GCP, error) { case "GKE": collector = compute.New(&compute.Config{ Projects: config.Projects, - }, computeService, billingService) + }, computeService, cloudCatalogClient) default: log.Printf("Unknown service %s", service) // Continue to next service, no need to halt here diff --git a/pkg/google/gcs/bucket_test.go b/pkg/google/gcs/bucket_test.go index 12d1e55c..b914cc2f 100644 --- a/pkg/google/gcs/bucket_test.go +++ b/pkg/google/gcs/bucket_test.go @@ -18,7 +18,7 @@ func TestNewBucketClient(t *testing.T) { tests := map[string]struct { client StorageClientInterface }{ - "Empty client": { + "Empty cloudCatalogClient": { client: gcs.NewStorageClientInterface(t), }, } @@ -26,7 +26,7 @@ func TestNewBucketClient(t *testing.T) { t.Run(name, func(t *testing.T) { client := NewBucketClient(test.client) if client == nil { - t.Errorf("expected client to be non-nil") + t.Errorf("expected cloudCatalogClient to be non-nil") } }) } diff --git a/pkg/google/gcs/gcs.go b/pkg/google/gcs/gcs.go index 2b4c5261..9624ac23 100644 --- a/pkg/google/gcs/gcs.go +++ b/pkg/google/gcs/gcs.go @@ -136,17 +136,17 @@ var operationsDiscountMap = map[string]map[string]map[string]float64{ } type Collector struct { - ProjectID string - Projects []string - client *billingv1.CloudCatalogClient - serviceName string - ctx context.Context - interval time.Duration - nextScrape time.Time - regionsClient RegionsClient - bucketClient *BucketClient - discount int - CachedBuckets *BucketCache + ProjectID string + Projects []string + cloudCatalogClient CloudCatalogClient + serviceName string + ctx context.Context + interval time.Duration + nextScrape time.Time + regionsClient RegionsClient + bucketClient *BucketClient + discount int + CachedBuckets *BucketCache } type Config struct { @@ -154,13 +154,19 @@ type Config struct { Projects string DefaultDiscount int ScrapeInterval time.Duration + ServiceName string } type RegionsClient interface { List(ctx context.Context, req *computepb.ListRegionsRequest, opts ...gax.CallOption) *compute.RegionIterator } -func New(config *Config, billingClient *billingv1.CloudCatalogClient, regionsClient RegionsClient, storageClient StorageClientInterface) (*Collector, error) { +type CloudCatalogClient interface { + ListServices(ctx context.Context, req *billingpb.ListServicesRequest, opts ...gax.CallOption) *billingv1.ServiceIterator + ListSkus(ctx context.Context, req *billingpb.ListSkusRequest, opts ...gax.CallOption) *billingv1.SkuIterator +} + +func New(config *Config, billingClient CloudCatalogClient, regionsClient RegionsClient, storageClient StorageClientInterface) (*Collector, error) { if config.ProjectId == "" { return nil, fmt.Errorf("projectID cannot be empty") } @@ -173,20 +179,16 @@ func New(config *Config, billingClient *billingv1.CloudCatalogClient, regionsCli } bucketClient := NewBucketClient(storageClient) - serviceName, err := getServiceNameByReadableName(ctx, billingClient, "Cloud Storage") - if err != nil { - return nil, err - } return &Collector{ - ProjectID: config.ProjectId, - Projects: projects, - client: billingClient, - regionsClient: regionsClient, - bucketClient: bucketClient, - discount: config.DefaultDiscount, - ctx: ctx, - serviceName: serviceName, - interval: config.ScrapeInterval, + ProjectID: config.ProjectId, + Projects: projects, + cloudCatalogClient: billingClient, + regionsClient: regionsClient, + bucketClient: bucketClient, + discount: config.DefaultDiscount, + ctx: ctx, + serviceName: config.ServiceName, + interval: config.ScrapeInterval, // Set nextScrape to the current time minus the scrape interval so that the first scrape will run immediately nextScrape: time.Now().Add(-config.ScrapeInterval), CachedBuckets: NewBucketCache(), @@ -197,7 +199,7 @@ func (c *Collector) Name() string { return "GCS" } -func getServiceNameByReadableName(ctx context.Context, client *billingv1.CloudCatalogClient, name string) (string, error) { +func GetServiceNameByReadableName(ctx context.Context, client CloudCatalogClient, name string) (string, error) { serviceList := client.ListServices(ctx, &billingpb.ListServicesRequest{}) for { row, err := serviceList.Next() @@ -249,7 +251,7 @@ func (c *Collector) Collect() error { if err != nil { log.Printf("Error exporting bucket info: %v", err) } - return ExportGCPCostData(c.ctx, c.client, c.serviceName) + return ExportGCPCostData(c.ctx, c.cloudCatalogClient, c.serviceName) } // ExportBucketInfo will list all buckets for a given project and export the data as a prometheus metric. @@ -330,7 +332,7 @@ func ExporterOperationsDiscounts() { } } -func ExportGCPCostData(ctx context.Context, client *billingv1.CloudCatalogClient, serviceName string) error { +func ExportGCPCostData(ctx context.Context, client CloudCatalogClient, serviceName string) error { skuResponse := client.ListSkus(ctx, &billingpb.ListSkusRequest{ Parent: serviceName, }) diff --git a/pkg/google/gcs/gcs_test.go b/pkg/google/gcs/gcs_test.go index 9c1333ef..0d068280 100644 --- a/pkg/google/gcs/gcs_test.go +++ b/pkg/google/gcs/gcs_test.go @@ -1,13 +1,10 @@ package gcs import ( - "context" "testing" - billingv1 "cloud.google.com/go/billing/apiv1" "cloud.google.com/go/billing/apiv1/billingpb" "github.com/stretchr/testify/assert" - "google.golang.org/api/option" "google.golang.org/genproto/googleapis/type/money" "github.com/grafana/cloudcost-exporter/mocks/pkg/google/gcs" @@ -139,12 +136,10 @@ func TestMisformedPricingInfoFromSku(t *testing.T) { } func TestNew(t *testing.T) { - billingClient, err := billingv1.NewCloudCatalogClient(context.Background(), option.WithAPIKey("hunter2")) - assert.NoError(t, err) + billingClient := gcs.NewCloudCatalogClient(t) regionsClient := gcs.NewRegionsClient(t) - assert.NoError(t, err) storageClient := gcs.NewStorageClientInterface(t) - t.Run("should return a non-nil client", func(t *testing.T) { + t.Run("should return a non-nil cloudCatalogClient", func(t *testing.T) { client, _ := New(&Config{ ProjectId: "project-1", }, billingClient, regionsClient, storageClient)