Refactor billing client

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,20 +65,25 @@ 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
 			}
 		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

pkg/google/gcs/bucket_test.go

@@ -18,15 +18,15 @@ func TestNewBucketClient(t *testing.T) {
 	tests := map[string]struct {
 		client StorageClientInterface
 	}{
-		"Empty client": {
+		"Empty cloudCatalogClient": {
 			client: gcs.NewStorageClientInterface(t),
 		},
 	}
 	for name, test := range tests {
 		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")
 			}
 		})
 	}

pkg/google/gcs/gcs.go

@@ -136,31 +136,37 @@ 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 {
 	ProjectId       string
 	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,
 	})

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)