fix: also uses the cluster name when getting the vpcid

topfreegames · Jun 22, 2023 · afc0870 · afc0870
1 parent 4f915be
commit afc0870
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 26 deletions.
diff --git a/controllers/clustermesh/clustermesh_controller.go b/controllers/clustermesh/clustermesh_controller.go
@@ -561,7 +561,7 @@ func PopulateClusterSpec(r *ClusterMeshReconciliation, ctx context.Context, clus
 		return clusterSpec, err
 	}
 
-	vpcId, err := ec2.GetVPCIdFromCIDR(ctx, r.ec2Client, r.kcp.Spec.KopsClusterSpec.NetworkCIDR)
+	vpcId, err := ec2.GetVPCIdWithCIDRAndClusterName(ctx, r.ec2Client, r.kcp.Name, r.kcp.Spec.KopsClusterSpec.NetworkCIDR)
 	if err != nil {
 		return clusterSpec, err
 	}

diff --git a/controllers/securitygroup/securitygroup_controller.go b/controllers/securitygroup/securitygroup_controller.go
@@ -117,7 +117,7 @@ func (r *SecurityGroupReconciliation) getAWSAccountInfo(ctx context.Context, kcp
 		return aws.String(""), aws.String(""), err
 	}
 
-	vpcId, err := ec2.GetVPCIdFromCIDR(ctx, r.ec2Client, kcp.Spec.KopsClusterSpec.NetworkCIDR)
+	vpcId, err := ec2.GetVPCIdWithCIDRAndClusterName(ctx, r.ec2Client, kcp.Name, kcp.Spec.KopsClusterSpec.NetworkCIDR)
 	if err != nil {
 		return aws.String(""), aws.String(""), err
 	}

diff --git a/pkg/aws/ec2/vpc.go b/pkg/aws/ec2/vpc.go
@@ -101,17 +101,22 @@ func GetReservationsUsingFilters(ctx context.Context, ec2Client EC2Client, filte
 	return instances.Reservations, nil
 }
 
-func GetVPCIdFromCIDR(ctx context.Context, ec2Client EC2Client, CIDR string) (*string, error) {
+func GetVPCIdWithCIDRAndClusterName(ctx context.Context, ec2Client EC2Client, clusterName, CIDR string) (*string, error) {
 
-	filter := "cidr"
 	result, err := ec2Client.DescribeVpcs(ctx, &ec2.DescribeVpcsInput{
 		Filters: []ec2types.Filter{
 			{
-				Name: &filter,
+				Name: aws.String("cidr"),
 				Values: []string{
 					CIDR,
 				},
 			},
+			{
+				Name: aws.String("tag:KubernetesCluster"),
+				Values: []string{
+					clusterName,
+				},
+			},
 		},
 	})
 	if err != nil {

diff --git a/pkg/aws/ec2/vpc_test.go b/pkg/aws/ec2/vpc_test.go
@@ -17,56 +17,73 @@ import (
 	"github.com/topfreegames/provider-crossplane/pkg/aws/ec2/fake"
 )
 
-func TestGetVPCIdFromCIDR(t *testing.T) {
-	testCases := []map[string]interface{}{
+func TestGetVPCIdWithCIDRAndClusterName(t *testing.T) {
+	testCases := []struct {
+		description      string
+		mockDescribeVpcs func(ctx context.Context, input *ec2.DescribeVpcsInput, opts []func(*ec2.Options)) (*ec2.DescribeVpcsOutput, error)
+		expectedError    error
+		expectedCIDR     string
+	}{
 		{
-			"description": "should return a VPC ID",
-			"mockDescribeVpcs": func(ctx context.Context, input *ec2.DescribeVpcsInput, opts []func(*ec2.Options)) (*ec2.DescribeVpcsOutput, error) {
+			description: "should return vpc-xxxxx",
+			mockDescribeVpcs: func(ctx context.Context, input *ec2.DescribeVpcsInput, opts []func(*ec2.Options)) (*ec2.DescribeVpcsOutput, error) {
 				return &ec2.DescribeVpcsOutput{
 					Vpcs: []ec2types.Vpc{
 						{
 							VpcId: aws.String("vpc-xxxxx"),
+							Tags: []ec2types.Tag{
+								{
+									Key:   aws.String("KubernetesCluster"),
+									Value: aws.String("test-cluster"),
+								},
+							},
+						},
+						{
+							VpcId: aws.String("vpc-yyyyy"),
+							Tags: []ec2types.Tag{
+								{
+									Key:   aws.String("KubernetesCluster"),
+									Value: aws.String("test-cluster-2"),
+								},
+							},
 						},
 					},
 				}, nil
 			},
-			"expectedError": false,
+			expectedCIDR: "vpc-xxxxx",
 		},
 		{
-			"description": "should fail to describe VPCs",
-			"mockDescribeVpcs": func(ctx context.Context, input *ec2.DescribeVpcsInput, opts []func(*ec2.Options)) (*ec2.DescribeVpcsOutput, error) {
+			description: "should fail to describe VPCs",
+			mockDescribeVpcs: func(ctx context.Context, input *ec2.DescribeVpcsInput, opts []func(*ec2.Options)) (*ec2.DescribeVpcsOutput, error) {
 				return nil, errors.New("some error")
 			},
-			"expectedError":        true,
-			"expectedErrorMessage": "failed to describe VPCs",
+			expectedError: fmt.Errorf("failed to describe VPCs"),
 		},
 		{
-			"description": "should fail to with empty result",
-			"mockDescribeVpcs": func(ctx context.Context, input *ec2.DescribeVpcsInput, opts []func(*ec2.Options)) (*ec2.DescribeVpcsOutput, error) {
+			description: "should fail with empty result",
+			mockDescribeVpcs: func(ctx context.Context, input *ec2.DescribeVpcsInput, opts []func(*ec2.Options)) (*ec2.DescribeVpcsOutput, error) {
 				return &ec2.DescribeVpcsOutput{
 					Vpcs: []ec2types.Vpc{},
 				}, nil
 			},
-			"expectedError":        true,
-			"expectedErrorMessage": "failed to retrieve vpc with CIDR",
+			expectedError: fmt.Errorf("failed to retrieve vpc with CIDR"),
 		},
 	}
 	RegisterFailHandler(Fail)
 	g := NewWithT(t)
 
 	for _, tc := range testCases {
-		t.Run(tc["description"].(string), func(t *testing.T) {
+		t.Run(tc.description, func(t *testing.T) {
 			ctx := context.TODO()
 			fakeEC2Client := &fake.MockEC2Client{}
-			fakeEC2Client.MockDescribeVpcs = tc["mockDescribeVpcs"].(func(ctx context.Context, input *ec2.DescribeVpcsInput, opts []func(*ec2.Options)) (*ec2.DescribeVpcsOutput, error))
-			vpcId, err := GetVPCIdFromCIDR(ctx, fakeEC2Client, "x.x.x.x/20")
+			fakeEC2Client.MockDescribeVpcs = tc.mockDescribeVpcs
+			vpcId, err := GetVPCIdWithCIDRAndClusterName(ctx, fakeEC2Client, "test-cluster", "x.x.x.x/20")
 
-			if !tc["expectedError"].(bool) {
-				g.Expect(err).To(BeNil())
-				g.Expect(vpcId).ToNot(BeNil())
+			if tc.expectedError != nil {
+				g.Expect(err.Error()).To(ContainSubstring(tc.expectedError.Error()))
 			} else {
-				g.Expect(err).ToNot(BeNil())
-				g.Expect(err.Error()).To(ContainSubstring(tc["expectedErrorMessage"].(string)))
+				g.Expect(err).To(BeNil())
+				g.Expect(*vpcId).To(BeEquivalentTo(tc.expectedCIDR))
 			}
 		})
 	}