[AWS ELB] Use correct region when fetching elb price locally (#915)

Co-authored-by: Stephan Rayner <stephan.rayner@grafana.com>

Author: leonorfmartins

pkg/aws/aws.go

@@ -208,8 +208,18 @@ func newWithDependencies(ctx context.Context, config *Config, awsClient client.C
 			}
 			collectors = append(collectors, collector)
 		case serviceELB:
+			// pricing API for ELB client needs to use always the same region
+			// as the pricing data is only available in us-east-1
+			elbPricingConfig, err := createAWSConfig(ctx, "us-east-1", config.Profile, config.RoleARN)
+			if err != nil {
+				return nil, err
+			}
+			awsELBPricingClient := client.NewAWSClient(client.Config{
+				PricingService: awsPricing.NewFromConfig(elbPricingConfig),
+			})
 			collector := elb.New(&elb.Config{
 				Regions:        regions,
+				PricingClient:  awsELBPricingClient,
 				RegionClients:  regionClients,
 				ScrapeInterval: config.ScrapeInterval,
 				Logger:         logger,

pkg/aws/elb/elb.go

@@ -19,7 +19,8 @@ import (
 )
 
 const (
-	subsystem = "aws_elb"
+	serviceName = "elb"
+	subsystem   = "aws_" + serviceName
 )
 
 var (
@@ -43,6 +44,7 @@ type Collector struct {
 	regions            []ec2Types.Region
 	ScrapeInterval     time.Duration
 	pricingMap         *ELBPricingMap
+	pricingClient      client.Client
 	awsRegionClientMap map[string]client.Client
 	NextScrape         time.Time
 	logger             *slog.Logger
@@ -52,9 +54,12 @@ type Collector struct {
 type Config struct {
 	ScrapeInterval time.Duration
 	Regions        []ec2Types.Region
-	RegionClients  map[string]client.Client
-	Logger         *slog.Logger
-	AccountID      string
+	// PricingClient must be a client configured for us-east-1: the AWS Pricing API
+	// is only available in us-east-1 and ap-south-1.
+	PricingClient client.Client
+	RegionClients map[string]client.Client
+	Logger        *slog.Logger
+	AccountID     string
 }
 
 type LoadBalancerInfo struct {
@@ -83,12 +88,14 @@ type elbProduct struct {
 }
 
 func New(config *Config) *Collector {
+	logger := config.Logger.With("logger", serviceName)
 	return &Collector{
 		regions:            config.Regions,
 		ScrapeInterval:     config.ScrapeInterval,
+		pricingClient:      config.PricingClient,
 		awsRegionClientMap: config.RegionClients,
-		logger:             config.Logger,
-		pricingMap:         NewELBPricingMap(config.Logger),
+		logger:             logger,
+		pricingMap:         NewELBPricingMap(logger),
 		accountID:          config.AccountID,
 	}
 }
@@ -107,7 +114,7 @@ func (c *Collector) Collect(ctx context.Context, ch chan<- prometheus.Metric) er
 	c.logger.Info("Starting ELB collection")
 
 	if c.shouldScrape() {
-		if err := c.pricingMap.refresh(ctx, c.awsRegionClientMap, c.regions); err != nil {
+		if err := c.pricingMap.refresh(ctx, c.pricingClient, c.regions); err != nil {
 			c.logger.Error("Failed to refresh pricing", "error", err)
 			return err
 		}

pkg/aws/elb/elb_test.go

@@ -1,6 +1,7 @@
 package elb
 
 import (
+	"errors"
 	"log/slog"
 	"testing"
 	"time"
@@ -24,6 +25,7 @@ func TestNew(t *testing.T) {
 		Regions: []ec2Types.Region{
 			{RegionName: stringPtr("us-east-1")},
 		},
+		PricingClient: mockClient,
 		RegionClients: map[string]client.Client{
 			"us-east-1": mockClient,
 		},
@@ -36,6 +38,7 @@ func TestNew(t *testing.T) {
 	assert.NotNil(t, collector)
 	assert.Equal(t, config.ScrapeInterval, collector.ScrapeInterval)
 	assert.Equal(t, config.Regions, collector.regions)
+	assert.Equal(t, mockClient, collector.pricingClient)
 	assert.Equal(t, mockClient, collector.awsRegionClientMap["us-east-1"])
 	assert.NotNil(t, collector.pricingMap)
 }
@@ -130,6 +133,58 @@ func TestCollectRegionLoadBalancers(t *testing.T) {
 
 }
 
+func TestFetchRegionPricing(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mockClient := mock_client.NewMockClient(ctrl)
+
+	albProduct := `{"Product":{"Attributes":{"usageType":"USE1-LoadBalancerUsage","operation":"LoadBalancing:Application"}},"Terms":{"OnDemand":{"t1":{"PriceDimensions":{"d1":{"pricePerUnit":{"USD":"0.0225"}}}}}}}`
+	nlbProduct := `{"Product":{"Attributes":{"usageType":"USE1-LCUUsage","operation":"LoadBalancing:Network"}},"Terms":{"OnDemand":{"t1":{"PriceDimensions":{"d1":{"pricePerUnit":{"USD":"0.006"}}}}}}}`
+	mockClient.EXPECT().ListELBPrices(gomock.Any(), "us-east-1").Return([]string{albProduct, nlbProduct}, nil)
+
+	pm := NewELBPricingMap(slog.Default())
+	pricing, err := pm.FetchRegionPricing(mockClient, t.Context(), "us-east-1")
+
+	assert.NoError(t, err)
+	assert.Equal(t, 0.0225, pricing.ALBHourlyRate[LoadBalancerUsage])
+	assert.Equal(t, 0.006, pricing.NLBHourlyRate[LCUUsage])
+}
+
+func TestFetchRegionPricingError(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mockClient := mock_client.NewMockClient(ctrl)
+	mockClient.EXPECT().ListELBPrices(gomock.Any(), "us-east-1").Return(nil, errors.New("api error"))
+
+	pm := NewELBPricingMap(slog.Default())
+	pricing, err := pm.FetchRegionPricing(mockClient, t.Context(), "us-east-1")
+
+	assert.Error(t, err)
+	assert.Nil(t, pricing)
+}
+
+func TestRefresh(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mockClient := mock_client.NewMockClient(ctrl)
+
+	albProduct := `{"Product":{"Attributes":{"usageType":"USE1-LoadBalancerUsage","operation":"LoadBalancing:Application"}},"Terms":{"OnDemand":{"t1":{"PriceDimensions":{"d1":{"pricePerUnit":{"USD":"0.0225"}}}}}}}`
+	mockClient.EXPECT().ListELBPrices(gomock.Any(), "us-east-1").Return([]string{albProduct}, nil)
+	mockClient.EXPECT().ListELBPrices(gomock.Any(), "us-west-2").Return([]string{albProduct}, nil)
+
+	pm := NewELBPricingMap(slog.Default())
+	regions := []ec2Types.Region{
+		{RegionName: stringPtr("us-east-1")},
+		{RegionName: stringPtr("us-west-2")},
+	}
+
+	err := pm.refresh(t.Context(), mockClient, regions)
+	assert.NoError(t, err)
+
+	for _, region := range []string{"us-east-1", "us-west-2"} {
+		pricing, err := pm.GetRegionPricing(region)
+		assert.NoError(t, err)
+		assert.Equal(t, 0.0225, pricing.ALBHourlyRate[LoadBalancerUsage])
+	}
+}
+
 func stringPtr(s string) *string {
 	return &s
 }

pkg/aws/elb/pricing_map.go

@@ -64,7 +64,7 @@ func (pm *ELBPricingMap) GetRegionPricing(region string) (*RegionPricing, error)
 	return pricing, nil
 }
 
-func (pm *ELBPricingMap) refresh(ctx context.Context, client map[string]client.Client, regions []ec2Types.Region) error {
+func (pm *ELBPricingMap) refresh(ctx context.Context, pricingClient client.Client, regions []ec2Types.Region) error {
 	pm.logger.Info("Refreshing ELB pricing data")
 
 	eg := errgroup.Group{}
@@ -73,7 +73,7 @@ func (pm *ELBPricingMap) refresh(ctx context.Context, client map[string]client.C
 	for _, region := range regions {
 		regionName := *region.RegionName
 		eg.Go(func() error {
-			pricing, err := pm.FetchRegionPricing(client[regionName], ctx, regionName)
+			pricing, err := pm.FetchRegionPricing(pricingClient, ctx, regionName)
 			if err != nil {
 				return fmt.Errorf("failed to fetch pricing for region %s: %w", regionName, err)
 			}