test: pkg/aws/s3

pkg/aws/costexplorer/costexplorer.go

@@ -0,0 +1,23 @@
+package costexplorer
+
+import (
+	"context"
+
+	"github.com/aws/aws-sdk-go-v2/service/costexplorer"
+)
+
+var (
+	_ CostExplorable = Client{}
+)
+
+type CostExplorable interface {
+	GetCostAndUsage(ctx context.Context, params *costexplorer.GetCostAndUsageInput, optFns ...func(*costexplorer.Options)) (*costexplorer.GetCostAndUsageOutput, error)
+}
+
+type Client struct {
+	c CostExplorable
+}
+
+func (c Client) GetCostAndUsage(ctx context.Context, params *costexplorer.GetCostAndUsageInput, optFns ...func(*costexplorer.Options)) (*costexplorer.GetCostAndUsageOutput, error) {
+	return c.GetCostAndUsage(ctx, params, optFns...)
+}

pkg/aws/s3/s3.go

@@ -9,10 +9,11 @@ import (
 	"time"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
-	"github.com/aws/aws-sdk-go-v2/service/costexplorer"
+	awscostexplorer "github.com/aws/aws-sdk-go-v2/service/costexplorer"
 	"github.com/aws/aws-sdk-go-v2/service/costexplorer/types"
 	"github.com/prometheus/client_golang/prometheus"
 
+	"github.com/grafana/cloudcost-exporter/pkg/aws/costexplorer"
 	"github.com/grafana/cloudcost-exporter/pkg/provider"
 )
 
@@ -103,13 +104,13 @@ var (
 // Collector is the AWS implementation of the Collector interface
 // It is responsible for registering and collecting metrics
 type Collector struct {
-	client     *costexplorer.Client
+	client     costexplorer.CostExplorable
 	interval   time.Duration
 	nextScrape time.Time
 }
 
 // New creates a new Collector with a client and scrape interval defined.
-func New(scrapeInterval time.Duration, client *costexplorer.Client) (*Collector, error) {
+func New(scrapeInterval time.Duration, client costexplorer.CostExplorable) (*Collector, error) {
 	return &Collector{
 		client:   client,
 		interval: scrapeInterval,
@@ -219,9 +220,9 @@ func (s S3BillingData) AddMetricGroup(region string, component string, group typ
 
 // getBillingData is responsible for making the API call to the AWS Cost Explorer API and parsing the response
 // into a S3BillingData struct
-func getBillingData(client *costexplorer.Client, startDate time.Time, endDate time.Time) (S3BillingData, error) {
+func getBillingData(client costexplorer.CostExplorable, startDate time.Time, endDate time.Time) (S3BillingData, error) {
 	log.Printf("Getting billing data for %s to %s\n", startDate.Format("2006-01-02"), endDate.Format("2006-01-02"))
-	input := &costexplorer.GetCostAndUsageInput{
+	input := &awscostexplorer.GetCostAndUsageInput{
 		TimePeriod: &types.DateInterval{
 			Start: aws.String(startDate.Format("2006-01-02")), // Specify the start date
 			End:   aws.String(endDate.Format("2006-01-02")),   // Specify the end date
@@ -243,7 +244,7 @@ func getBillingData(client *costexplorer.Client, startDate time.Time, endDate ti
 		},
 	}
 
-	var outputs []*costexplorer.GetCostAndUsageOutput
+	var outputs []*awscostexplorer.GetCostAndUsageOutput
 	for {
 		RequestCount.Inc()
 		output, err := client.GetCostAndUsage(context.TODO(), input)
@@ -263,7 +264,7 @@ func getBillingData(client *costexplorer.Client, startDate time.Time, endDate ti
 }
 
 // parseBillingData takes the output from the AWS Cost Explorer API and parses it into a S3BillingData struct
-func parseBillingData(outputs []*costexplorer.GetCostAndUsageOutput) S3BillingData {
+func parseBillingData(outputs []*awscostexplorer.GetCostAndUsageOutput) S3BillingData {
 	billingData := NewS3BillingData()
 
 	// Process the billing data in the 'output' variable
@@ -328,7 +329,7 @@ func getComponentFromKey(key string) string {
 }
 
 // ExportBillingData will query the previous 30 days of S3 billing data and export it to the prometheus metrics
-func ExportBillingData(client *costexplorer.Client) error {
+func ExportBillingData(client costexplorer.CostExplorable) error {
 	// We go one day into the past as the current days billing data has no guarantee of being complete
 	endDate := time.Now().AddDate(0, 0, -1)
 	// Current assumption is that we're going to pull 30 days worth of billing data

pkg/aws/s3/s3_test.go

@@ -1,16 +1,23 @@
 package s3
 
 import (
+	"context"
 	"encoding/csv"
+	"fmt"
 	"os"
 	"testing"
 	"time"
 
-	"github.com/aws/aws-sdk-go-v2/service/costexplorer"
+	awscostexplorer "github.com/aws/aws-sdk-go-v2/service/costexplorer"
+	"github.com/aws/aws-sdk-go-v2/service/costexplorer/types"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/mock/gomock"
 
-	"github.com/aws/aws-sdk-go-v2/service/costexplorer/types"
+	mockcostexplorer "github.com/grafana/cloudcost-exporter/mocks/pkg/aws/costexplorer"
+	"github.com/grafana/cloudcost-exporter/pkg/aws/costexplorer"
+	mock_provider "github.com/grafana/cloudcost-exporter/pkg/provider/mocks"
 )
 
 func Test_getDimensionFromKey(t *testing.T) {
@@ -153,7 +160,7 @@ func TestS3BillingData_AddRegion(t *testing.T) {
 func TestNewCollector(t *testing.T) {
 	type args struct {
 		interval time.Duration
-		client   *costexplorer.Client
+		client   costexplorer.CostExplorable
 	}
 	tests := map[string]struct {
 		args  args
@@ -170,7 +177,9 @@ func TestNewCollector(t *testing.T) {
 	}
 	for name, tt := range tests {
 		t.Run(name, func(t *testing.T) {
-			got, err := New(tt.args.interval, tt.args.client)
+			c := mockcostexplorer.NewCostExplorable(t)
+
+			got, err := New(tt.args.interval, c)
 			if tt.error {
 				require.Error(t, err)
 				return
@@ -181,3 +190,51 @@ func TestNewCollector(t *testing.T) {
 		})
 	}
 }
+
+func TestCollector_Name(t *testing.T) {
+	c := &Collector{}
+	require.Equal(t, "S3", c.Name())
+}
+
+func TestCollector_Register(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	r := mock_provider.NewMockRegistry(ctrl)
+	r.EXPECT().MustRegister(gomock.Any()).Times(5)
+
+	c := &Collector{}
+	err := c.Register(r)
+	require.NoError(t, err)
+}
+
+func TestExportBillingData(t *testing.T) {
+	for _, tc := range []struct {
+		name            string
+		GetCostAndUsage func(ctx context.Context, params *awscostexplorer.GetCostAndUsageInput, optFns ...func(*awscostexplorer.Options)) (*awscostexplorer.GetCostAndUsageOutput, error)
+		expectedError   error
+	}{
+		{
+			name: "error",
+			GetCostAndUsage: func(ctx context.Context, params *awscostexplorer.GetCostAndUsageInput, optFns ...func(*awscostexplorer.Options)) (*awscostexplorer.GetCostAndUsageOutput, error) {
+				return nil, fmt.Errorf("test cost and usage error")
+			},
+			expectedError: fmt.Errorf("test cost and usage error"),
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			ce := mockcostexplorer.NewCostExplorable(t)
+			if tc.GetCostAndUsage != nil {
+				ce.EXPECT().
+					GetCostAndUsage(mock.Anything, mock.Anything, mock.Anything).
+					RunAndReturn(tc.GetCostAndUsage).
+					Once()
+			}
+
+			err := ExportBillingData(ce)
+			if tc.expectedError != nil {
+				require.EqualError(t, err, tc.expectedError.Error())
+				return
+			}
+			require.NoError(t, err)
+		})
+	}
+}