From 597eb701ba7608f1c7f52e42168eb089d592ca6e Mon Sep 17 00:00:00 2001
From: phuhung273 <phuhung.tranpham@gmail.com>
Date: Fri, 20 Dec 2024 04:29:49 +0000
Subject: [PATCH] cleanup

---
 pkg/networking/backend_sg_provider.go      |  5 +-
 pkg/networking/backend_sg_provider_test.go | 96 +++++++++++++---------
 2 files changed, 58 insertions(+), 43 deletions(-)

diff --git a/pkg/networking/backend_sg_provider.go b/pkg/networking/backend_sg_provider.go
index 1bb141105..6f8f1974b 100644
--- a/pkg/networking/backend_sg_provider.go
+++ b/pkg/networking/backend_sg_provider.go
@@ -5,15 +5,14 @@ import (
 	"crypto/sha256"
 	"encoding/hex"
 	"fmt"
+	ec2types "github.com/aws/aws-sdk-go-v2/service/ec2/types"
+	"github.com/aws/smithy-go"
 	"regexp"
 	"sort"
 	"strings"
 	"sync"
 	"time"
 
-	ec2types "github.com/aws/aws-sdk-go-v2/service/ec2/types"
-	"github.com/aws/smithy-go"
-
 	awssdk "github.com/aws/aws-sdk-go-v2/aws"
 	ec2sdk "github.com/aws/aws-sdk-go-v2/service/ec2"
 	"github.com/go-logr/logr"
diff --git a/pkg/networking/backend_sg_provider_test.go b/pkg/networking/backend_sg_provider_test.go
index 81b79bd5b..bcd551e96 100644
--- a/pkg/networking/backend_sg_provider_test.go
+++ b/pkg/networking/backend_sg_provider_test.go
@@ -2,13 +2,12 @@ package networking
 
 import (
 	"context"
-	"reflect"
-	"testing"
-
 	ec2types "github.com/aws/aws-sdk-go-v2/service/ec2/types"
 	"github.com/aws/smithy-go"
 	"k8s.io/apimachinery/pkg/types"
+	"reflect"
 	"sigs.k8s.io/aws-load-balancer-controller/pkg/k8s"
+	"testing"
 
 	"github.com/go-logr/logr"
 	corev1 "k8s.io/api/core/v1"
@@ -48,13 +47,13 @@ func Test_defaultBackendSGProvider_Get(t *testing.T) {
 		err  error
 	}
 	type fields struct {
-		backendSG         string
-		ingResources      []*networking.Ingress
-		svcResource       *corev1.Service
-		defaultTags       map[string]string
-		describeSGCalls   []describeSecurityGroupsAsListCall
-		createSGCalls     []createSecurityGroupWithContexCall
-		createSGTagsCalls []createTagsWithContextCall
+		backendSG                  string
+		ingResources               []*networking.Ingress
+		svcResource                *corev1.Service
+		defaultTags                map[string]string
+		describeSGCalls            []describeSecurityGroupsAsListCall
+		createSGCalls              []createSecurityGroupWithContexCall
+		createTagsWithContextCalls []createTagsWithContextCall
 	}
 	defaultEC2Filters := []ec2types.Filter{
 		{
@@ -94,33 +93,50 @@ func Test_defaultBackendSGProvider_Get(t *testing.T) {
 		fields  fields
 		wantErr error
 	}{
-		// {
-		// 	name: "backend sg enabled",
-		// 	fields: fields{
-		// 		backendSG:    "sg-xxx",
-		// 		ingResources: []*networking.Ingress{ing},
-		// 	},
-		// 	want: "sg-xxx",
-		// },
-		// {
-		// 	name: "backend sg enabled, auto-gen, SG exists",
-		// 	fields: fields{
-		// 		describeSGCalls: []describeSecurityGroupsAsListCall{
-		// 			{
-		// 				req: &ec2sdk.DescribeSecurityGroupsInput{
-		// 					Filters: defaultEC2Filters,
-		// 				},
-		// 				resp: []ec2types.SecurityGroup{
-		// 					{
-		// 						GroupId: awssdk.String("sg-autogen"),
-		// 					},
-		// 				},
-		// 			},
-		// 		},
-		// 		ingResources: []*networking.Ingress{ing, ing1},
-		// 	},
-		// 	want: "sg-autogen",
-		// },
+		{
+			name: "backend sg enabled",
+			fields: fields{
+				backendSG:    "sg-xxx",
+				ingResources: []*networking.Ingress{ing},
+			},
+			want: "sg-xxx",
+		},
+		{
+			name: "backend sg enabled, auto-gen, SG exists",
+			fields: fields{
+				describeSGCalls: []describeSecurityGroupsAsListCall{
+					{
+						req: &ec2sdk.DescribeSecurityGroupsInput{
+							Filters: defaultEC2Filters,
+						},
+						resp: []ec2types.SecurityGroup{
+							{
+								GroupId: awssdk.String("sg-autogen"),
+							},
+						},
+					},
+				},
+				createTagsWithContextCalls: []createTagsWithContextCall{
+					{
+						req: &ec2sdk.CreateTagsInput{
+							Resources: []string{"sg-autogen"},
+							Tags: []ec2types.Tag{
+								{
+									Key:   awssdk.String("elbv2.k8s.aws/cluster"),
+									Value: awssdk.String(defaultClusterName),
+								},
+								{
+									Key:   awssdk.String("elbv2.k8s.aws/resource"),
+									Value: awssdk.String("backend-sg"),
+								},
+							},
+						},
+					},
+				},
+				ingResources: []*networking.Ingress{ing, ing1},
+			},
+			want: "sg-autogen",
+		},
 		{
 			name: "backend sg enabled, auto-gen, SG exists, try to sync tags",
 			fields: fields{
@@ -136,7 +152,7 @@ func Test_defaultBackendSGProvider_Get(t *testing.T) {
 						},
 					},
 				},
-				createSGTagsCalls: []createTagsWithContextCall{
+				createTagsWithContextCalls: []createTagsWithContextCall{
 					{
 						req: &ec2sdk.CreateTagsInput{
 							Resources: []string{"sg-autogen"},
@@ -189,7 +205,7 @@ func Test_defaultBackendSGProvider_Get(t *testing.T) {
 						},
 					},
 				},
-				createSGTagsCalls: []createTagsWithContextCall{
+				createTagsWithContextCalls: []createTagsWithContextCall{
 					{
 						req: &ec2sdk.CreateTagsInput{
 							Resources: []string{"sg-autogen"},
@@ -397,7 +413,7 @@ func Test_defaultBackendSGProvider_Get(t *testing.T) {
 			for _, call := range tt.fields.createSGCalls {
 				ec2Client.EXPECT().CreateSecurityGroupWithContext(context.Background(), call.req).Return(call.resp, call.err)
 			}
-			for _, call := range tt.fields.createSGTagsCalls {
+			for _, call := range tt.fields.createTagsWithContextCalls {
 				ec2Client.EXPECT().CreateTagsWithContext(context.Background(), call.req).Return(call.resp, call.err)
 			}
 			k8sClient := mock_client.NewMockClient(ctrl)