From eaee2eeed15bc992d7b3c292eb721f19a5621f42 Mon Sep 17 00:00:00 2001
From: zhewang <zwan@redhat.com>
Date: Fri, 10 Jan 2025 12:25:18 +0800
Subject: [PATCH] OCM-13401 | test: Update DestroyBastionProxy func.

---
 pkg/aws/aws_client/instance.go |  1 +
 pkg/test/vpc_client/bastion.go | 61 ++++++++++++++++++++++++++++++----
 2 files changed, 56 insertions(+), 6 deletions(-)

diff --git a/pkg/aws/aws_client/instance.go b/pkg/aws/aws_client/instance.go
index b6462e2..38e285d 100644
--- a/pkg/aws/aws_client/instance.go
+++ b/pkg/aws/aws_client/instance.go
@@ -83,6 +83,7 @@ func (client *AWSClient) WaitForInstanceReady(instanceID string, timeout time.Du
 		instanceID,
 	}
 	log.LogInfo("Waiting for below instances ready: %s ", strings.Join(instanceIDs, "|"))
+	time.Sleep(2 * time.Second)
 	_, err := client.WaitForInstancesRunning(instanceIDs, 10)
 	return err
 }
diff --git a/pkg/test/vpc_client/bastion.go b/pkg/test/vpc_client/bastion.go
index 27b4641..da8862f 100644
--- a/pkg/test/vpc_client/bastion.go
+++ b/pkg/test/vpc_client/bastion.go
@@ -8,6 +8,7 @@ import (
 	"github.com/openshift-online/ocm-common/pkg/file"
 	"golang.org/x/crypto/bcrypt"
 	"net/url"
+	"strings"
 	"time"
 
 	CON "github.com/openshift-online/ocm-common/pkg/aws/consts"
@@ -130,15 +131,63 @@ func (vpc *VPC) PrepareBastionProxy(zone string, keypairName string, privateKeyP
 	return proxyUrl, nil
 }
 
-func (vpc *VPC) DestroyBastionProxy(instance types.Instance) error {
-	var instanceIDs []string
-	instanceIDs = append(instanceIDs, *instance.InstanceId)
-	err := vpc.AWSClient.TerminateInstances(instanceIDs, true, 10)
+func (vpc *VPC) DestroyBastionProxy() error {
+	filters := []map[string][]string{
+		{
+			"vpc-id": []string{
+				vpc.VpcID,
+			},
+		},
+	}
+	filters = append(filters, map[string][]string{
+		"tag:Name": {
+			CON.BastionName,
+		},
+	})
+
+	insts, err := vpc.AWSClient.ListInstances([]string{}, filters...)
+
+	if err != nil {
+		log.LogError("Error happened when list instances for vpc %s: %s", vpc.VpcID, err)
+		return err
+	}
+	needTermination := []string{}
+	keyPairNames := []string{}
+	for _, inst := range insts {
+		needTermination = append(needTermination, *inst.InstanceId)
+		if inst.KeyName != nil {
+			keyPairNames = append(keyPairNames, *inst.KeyName)
+		}
+	}
+	err = vpc.AWSClient.TerminateInstances(needTermination, true, 20)
+	if err != nil {
+		log.LogError("Terminating instances %s meet error: %s", strings.Join(needTermination, ","), err)
+	} else {
+		log.LogInfo("Terminating instances %s successfully", strings.Join(needTermination, ","))
+	}
+	err = vpc.DeleteKeyPair(keyPairNames)
+	if err != nil {
+		log.LogError("Delete key pair %s meet error: %s", strings.Join(keyPairNames, ","), err)
+	}
+	needCleanGroups := []types.SecurityGroup{}
+	securityGroups, err := vpc.AWSClient.ListSecurityGroups(vpc.VpcID)
 	if err != nil {
-		log.LogError("Terminate instance failed")
 		return err
 	}
-	return nil
+	for _, sg := range securityGroups {
+		for _, tag := range sg.Tags {
+			if *tag.Key == "Name" && *tag.Value == CON.AdditionalSecurityGroupName {
+				needCleanGroups = append(needCleanGroups, sg)
+			}
+		}
+	}
+	for _, sg := range needCleanGroups {
+		_, err = vpc.AWSClient.DeleteSecurityGroup(*sg.GroupId)
+		if err != nil {
+			return err
+		}
+	}
+	return err
 }
 
 func generateBcryptPassword(plainPassword string) (string, error) {