From 39dfbe79511b5bab48c9659f91c85e125f5a1049 Mon Sep 17 00:00:00 2001
From: Marco Dinis <marco.dinis@goteleport.com>
Date: Tue, 10 Dec 2024 11:33:22 +0000
Subject: [PATCH] extract loops

---
 lib/web/integrations_awsoidc.go | 91 ++++++++++++++++++++-------------
 lib/web/ui/integration.go       |  2 +-
 2 files changed, 57 insertions(+), 36 deletions(-)

diff --git a/lib/web/integrations_awsoidc.go b/lib/web/integrations_awsoidc.go
index 79e5b295eae8f..bfb7874b9de60 100644
--- a/lib/web/integrations_awsoidc.go
+++ b/lib/web/integrations_awsoidc.go
@@ -24,6 +24,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"log/slog"
+	"maps"
 	"net/http"
 	"net/url"
 	"slices"
@@ -287,15 +288,16 @@ func (h *Handler) awsOIDCListDeployedDatabaseService(w http.ResponseWriter, r *h
 	}, nil
 }
 
-func fetchRelevantAWSRegions(ctx context.Context,
-	authClient interface {
-		GetResources(ctx context.Context, req *proto.ListResourcesRequest) (*proto.ListResourcesResponse, error)
-		GetDatabases(context.Context) ([]types.Database, error)
-	},
-	discoveryConfigsClient interface {
-		ListDiscoveryConfigs(ctx context.Context, pageSize int, nextToken string) ([]*discoveryconfig.DiscoveryConfig, string, error)
-	},
-) ([]string, error) {
+type databaseGetter interface {
+	GetResources(ctx context.Context, req *proto.ListResourcesRequest) (*proto.ListResourcesResponse, error)
+	GetDatabases(context.Context) ([]types.Database, error)
+}
+
+type discoveryConfigLister interface {
+	ListDiscoveryConfigs(ctx context.Context, pageSize int, nextToken string) ([]*discoveryconfig.DiscoveryConfig, string, error)
+}
+
+func fetchRelevantAWSRegions(ctx context.Context, authClient databaseGetter, discoveryConfigsClient discoveryConfigLister) ([]string, error) {
 	regionsSet := make(map[string]struct{})
 
 	// Collect Regions from Database resources.
@@ -321,21 +323,9 @@ func fetchRelevantAWSRegions(ctx context.Context,
 		if err != nil {
 			return nil, trace.Wrap(err)
 		}
-		for _, resource := range page.Resources {
-			for _, matcher := range resource.GetResourceMatchers() {
-				if matcher.Labels == nil {
-					continue
-				}
-				for labelKey, labelValues := range *matcher.Labels {
-					if labelKey != types.DiscoveryLabelRegion {
-						continue
-					}
-					for _, labelValue := range labelValues {
-						regionsSet[labelValue] = struct{}{}
-					}
-				}
-			}
-		}
+
+		maps.Copy(regionsSet, extractRegionsFromDatabaseServicesPage(page.Resources))
+
 		if page.NextKey == "" {
 			break
 		}
@@ -350,13 +340,7 @@ func fetchRelevantAWSRegions(ctx context.Context,
 			return nil, trace.Wrap(err)
 		}
 
-		for _, dc := range resp {
-			for _, awsMatcher := range dc.Spec.AWS {
-				for _, region := range awsMatcher.Regions {
-					regionsSet[region] = struct{}{}
-				}
-			}
-		}
+		maps.Copy(regionsSet, extractRegionsFromDiscoveryConfigPage(resp))
 
 		if respNextPageKey == "" {
 			break
@@ -375,19 +359,56 @@ func fetchRelevantAWSRegions(ctx context.Context,
 	return ret, nil
 }
 
+func extractRegionsFromDatabaseServicesPage(dbServices []types.DatabaseService) map[string]struct{} {
+	regionsSet := make(map[string]struct{})
+	for _, resource := range dbServices {
+		for _, matcher := range resource.GetResourceMatchers() {
+			if matcher.Labels == nil {
+				continue
+			}
+			for labelKey, labelValues := range *matcher.Labels {
+				if labelKey != types.DiscoveryLabelRegion {
+					continue
+				}
+				for _, labelValue := range labelValues {
+					regionsSet[labelValue] = struct{}{}
+				}
+			}
+		}
+	}
+
+	return regionsSet
+}
+
+func extractRegionsFromDiscoveryConfigPage(discoveryConfigs []*discoveryconfig.DiscoveryConfig) map[string]struct{} {
+	regionsSet := make(map[string]struct{})
+
+	for _, dc := range discoveryConfigs {
+		for _, awsMatcher := range dc.Spec.AWS {
+			for _, region := range awsMatcher.Regions {
+				regionsSet[region] = struct{}{}
+			}
+		}
+	}
+
+	return regionsSet
+}
+
+type deployedDatabaseServiceLister interface {
+	ListDeployedDatabaseServices(ctx context.Context, in *integrationv1.ListDeployedDatabaseServicesRequest, opts ...grpc.CallOption) (*integrationv1.ListDeployedDatabaseServicesResponse, error)
+}
+
 func listDeployedDatabaseServices(ctx context.Context,
 	logger *slog.Logger,
 	integrationName string,
 	regions []string,
-	awsoidcClient interface {
-		ListDeployedDatabaseServices(ctx context.Context, in *integrationv1.ListDeployedDatabaseServicesRequest, opts ...grpc.CallOption) (*integrationv1.ListDeployedDatabaseServicesResponse, error)
-	},
+	awsOIDCClient deployedDatabaseServiceLister,
 ) ([]ui.AWSOIDCDeployedDatabaseService, error) {
 	var services []ui.AWSOIDCDeployedDatabaseService
 	for _, region := range regions {
 		var nextToken string
 		for {
-			resp, err := awsoidcClient.ListDeployedDatabaseServices(ctx, &integrationv1.ListDeployedDatabaseServicesRequest{
+			resp, err := awsOIDCClient.ListDeployedDatabaseServices(ctx, &integrationv1.ListDeployedDatabaseServicesRequest{
 				Integration: integrationName,
 				Region:      region,
 				NextToken:   nextToken,
diff --git a/lib/web/ui/integration.go b/lib/web/ui/integration.go
index 0a616a4ba8947..63369beacf95c 100644
--- a/lib/web/ui/integration.go
+++ b/lib/web/ui/integration.go
@@ -371,7 +371,7 @@ type AWSOIDCDeployDatabaseServiceResponse struct {
 	ClusterDashboardURL string `json:"clusterDashboardUrl"`
 }
 
-// AWSOIDCDeployedDatabaseService represents an Teleport Database Service that is deployed in Amazon ECS.
+// AWSOIDCDeployedDatabaseService represents a Teleport Database Service that is deployed in Amazon ECS.
 type AWSOIDCDeployedDatabaseService struct {
 	// Name is the ECS Service name.
 	Name string `json:"name,omitempty"`