From 166bd133bff87106f96eec47d13a21562a7c5a97 Mon Sep 17 00:00:00 2001
From: Chao Chen <chaochn@amazon.com>
Date: Tue, 4 Apr 2023 00:06:30 -0700
Subject: [PATCH] tests/common: kill etcd cluster processes after test finished
 only after peerTLS is turned on

Signed-off-by: Chao Chen <chaochn@amazon.com>
---
 tests/common/e2e_test.go       |  4 ++--
 tests/framework/e2e/cluster.go | 21 +++++++++++++++++----
 tests/framework/e2e/config.go  |  3 ++-
 tests/framework/e2e/e2e.go     |  1 +
 4 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/tests/common/e2e_test.go b/tests/common/e2e_test.go
index 11c4f94a335b..b47bcf0da7eb 100644
--- a/tests/common/e2e_test.go
+++ b/tests/common/e2e_test.go
@@ -36,11 +36,11 @@ func e2eClusterTestCases() []testCase {
 		},
 		{
 			name:   "PeerTLS",
-			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.ManualTLS},
+			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.ManualTLS, ClusterContext: &e2e.ClusterContext{ForceStop: true}},
 		},
 		{
 			name:   "PeerAutoTLS",
-			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.AutoTLS},
+			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.AutoTLS, ClusterContext: &e2e.ClusterContext{ForceStop: true}},
 		},
 		{
 			name:   "ClientTLS",
diff --git a/tests/framework/e2e/cluster.go b/tests/framework/e2e/cluster.go
index 35e83bc2f229..f3a6b31a3d83 100644
--- a/tests/framework/e2e/cluster.go
+++ b/tests/framework/e2e/cluster.go
@@ -130,8 +130,9 @@ type EtcdProcessCluster struct {
 }
 
 type EtcdProcessClusterConfig struct {
-	Logger  *zap.Logger
-	Version ClusterVersion
+	Logger    *zap.Logger
+	Version   ClusterVersion
+	ForceStop bool
 	// DataDirPath specifies the data-dir for the members. If test cases
 	// do not specify `DataDirPath`, then e2e framework creates a
 	// temporary directory for each member; otherwise, it creates a
@@ -886,9 +887,14 @@ func (epc *EtcdProcessCluster) Etcdctl(opts ...config.ClientOption) *EtcdctlV3 {
 	return etcdctl
 }
 
-func (epc *EtcdProcessCluster) Close() error {
+func (epc *EtcdProcessCluster) Close() (err error) {
 	epc.lg.Info("closing test cluster...")
-	err := epc.Stop()
+	if epc.Cfg.ForceStop {
+		err = epc.ForceStop()
+	} else {
+		err = epc.Stop()
+	}
+
 	for _, p := range epc.Procs {
 		// p is nil when NewEtcdProcess fails in the middle
 		// Close still gets called to clean up test data
@@ -903,6 +909,13 @@ func (epc *EtcdProcessCluster) Close() error {
 	return err
 }
 
+func (epc *EtcdProcessCluster) ForceStop() (err error) {
+	for _, member := range epc.Procs {
+		member.Kill()
+	}
+	return epc.Stop()
+}
+
 func findMemberIDByEndpoint(members []*etcdserverpb.Member, endpoint string) (uint64, error) {
 	for _, m := range members {
 		if m.ClientURLs[0] == endpoint {
diff --git a/tests/framework/e2e/config.go b/tests/framework/e2e/config.go
index acc1d82e0484..5f33272d12e4 100644
--- a/tests/framework/e2e/config.go
+++ b/tests/framework/e2e/config.go
@@ -24,5 +24,6 @@ const (
 )
 
 type ClusterContext struct {
-	Version ClusterVersion
+	Version   ClusterVersion
+	ForceStop bool
 }
diff --git a/tests/framework/e2e/e2e.go b/tests/framework/e2e/e2e.go
index f78df57926ea..eae3dd650256 100644
--- a/tests/framework/e2e/e2e.go
+++ b/tests/framework/e2e/e2e.go
@@ -56,6 +56,7 @@ func (e e2eRunner) NewCluster(ctx context.Context, t testing.TB, opts ...config.
 	if cfg.ClusterContext != nil {
 		e2eClusterCtx := cfg.ClusterContext.(*ClusterContext)
 		e2eConfig.Version = e2eClusterCtx.Version
+		e2eConfig.ForceStop = e2eClusterCtx.ForceStop
 	}
 
 	switch cfg.ClientTLS {