From 847396655c8d64dce1eb053abb8ec335831f09df Mon Sep 17 00:00:00 2001
From: Ola Saadi <ola.saadi1@ibm.com>
Date: Thu, 21 Nov 2024 11:56:13 +0200
Subject: [PATCH] fix

---
 pkg/vpcmodel/subnetsConnectivity.go | 17 ++++++++++++++---
 pkg/vpcmodel/textOutput.go          | 18 ++++++++++++++----
 2 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/pkg/vpcmodel/subnetsConnectivity.go b/pkg/vpcmodel/subnetsConnectivity.go
index 5e3fccf97..1389e8d6b 100644
--- a/pkg/vpcmodel/subnetsConnectivity.go
+++ b/pkg/vpcmodel/subnetsConnectivity.go
@@ -241,10 +241,21 @@ func updateSubnetsConnectivityByTransitGateway(src, dst VPCResourceIntf,
 	c *VPCConfig) (
 	*netset.TransportSet, error) {
 	// assuming a single router representing the tgw for a "MultipleVPCsConfig"
-	if len(c.RoutingResources) != 2 { // expecting tgw and sgw (virtual gateway)
-		return nil, fmt.Errorf("unexpected number of RoutingResources for MultipleVPCsConfig, expecting only TGW")
+	resourceTypeTGW := "TGW"
+	var tgw RoutingResource
+	tgwRouterFound := false
+	for _, router := range c.RoutingResources {
+		if router.Kind() == resourceTypeTGW {
+			if tgwRouterFound {
+				return nil, fmt.Errorf("unexpected number of RoutingResources for MultipleVPCsConfig, expecting only one TGW")
+			}
+			tgw = router
+			tgwRouterFound = true
+		}
+	}
+	if !tgwRouterFound {
+		return nil, fmt.Errorf("unexpected number of RoutingResources for MultipleVPCsConfig, expecting TGW")
 	}
-	tgw := c.RoutingResources[0]
 	connections, err := tgw.AllowedConnectivity(src, dst)
 	if err != nil {
 		return nil, err
diff --git a/pkg/vpcmodel/textOutput.go b/pkg/vpcmodel/textOutput.go
index 887c41e48..080a00edf 100644
--- a/pkg/vpcmodel/textOutput.go
+++ b/pkg/vpcmodel/textOutput.go
@@ -15,11 +15,21 @@ type TextOutputFormatter struct {
 }
 
 func multipleVPCsConfigHeader(c *VPCConfig) (string, error) {
-	if len(c.RoutingResources) != 2 { // expecting tgw and sgw (virtual gateway)
-		fmt.Printf("c.RoutingResources: %v\n", c.RoutingResources)
-		return "", errors.New("unexpected config of multiple VPCs connected by TGW, missing TGW resource")
+	resourceTypeTGW := "TGW"
+	var tgw RoutingResource
+	tgwRouterFound := false
+	for _, router := range c.RoutingResources {
+		if router.Kind() == resourceTypeTGW {
+			if tgwRouterFound {
+				return "", errors.New("unexpected number of RoutingResources for MultipleVPCsConfig, expecting only one TGW")
+			}
+			tgw = router
+			tgwRouterFound = true
+		}
+	}
+	if !tgwRouterFound {
+		return "", errors.New("unexpected number of RoutingResources for MultipleVPCsConfig, expecting TGW")
 	}
-	tgw := c.RoutingResources[0]
 	return fmt.Sprintf("Connectivity between VPCs connected by TGW %s (UID: %s)\n", tgw.NameForAnalyzerOut(c), tgw.UID()), nil
 }