From 15efc559051a44073a55ca661c28d414dfb26ad0 Mon Sep 17 00:00:00 2001
From: Wei Fu <fuweid89@gmail.com>
Date: Fri, 7 Jul 2023 21:28:51 +0800
Subject: [PATCH] tests/e2e: allow to use SIGTERM to verify graceful-stop

Signed-off-by: Wei Fu <fuweid89@gmail.com>
---
 tests/e2e/cluster_test.go | 10 ++++++++++
 tests/e2e/cmux_test.go    | 14 ++++++++++++--
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/tests/e2e/cluster_test.go b/tests/e2e/cluster_test.go
index 9186accd181..f13fc88c590 100644
--- a/tests/e2e/cluster_test.go
+++ b/tests/e2e/cluster_test.go
@@ -133,6 +133,8 @@ type etcdProcessClusterConfig struct {
 	WatchProcessNotifyInterval time.Duration
 
 	debug bool
+
+	stopSignal os.Signal
 }
 
 // newEtcdProcessCluster launches a new cluster from etcd processes, returning
@@ -151,12 +153,20 @@ func newEtcdProcessCluster(cfg *etcdProcessClusterConfig) (*etcdProcessCluster,
 			epc.Close()
 			return nil, err
 		}
+
 		epc.procs[i] = proc
 	}
 
 	if err := epc.Start(); err != nil {
 		return nil, err
 	}
+
+	// overwrite the default signal
+	if cfg.stopSignal != nil {
+		for _, proc := range epc.procs {
+			proc.WithStopSignal(cfg.stopSignal)
+		}
+	}
 	return epc, nil
 }
 
diff --git a/tests/e2e/cmux_test.go b/tests/e2e/cmux_test.go
index df9ee3ac114..8d394b31c66 100644
--- a/tests/e2e/cmux_test.go
+++ b/tests/e2e/cmux_test.go
@@ -21,6 +21,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"syscall"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -63,10 +64,19 @@ func TestConnectionMultiplexing(t *testing.T) {
 	} {
 		t.Run(tc.name, func(t *testing.T) {
 			ctx := context.Background()
-			cfg := etcdProcessClusterConfig{clusterSize: 1, clientTLS: tc.serverTLS, enableV2: true, clientHttpSeparate: tc.separateHttpPort}
+			cfg := etcdProcessClusterConfig{
+				clusterSize:        1,
+				clientTLS:          tc.serverTLS,
+				enableV2:           true,
+				clientHttpSeparate: tc.separateHttpPort,
+				stopSignal:         syscall.SIGTERM, // check graceful stop
+			}
 			clus, err := newEtcdProcessCluster(&cfg)
 			require.NoError(t, err)
-			defer clus.Close()
+
+			defer func() {
+				require.NoError(t, clus.Close())
+			}()
 
 			var clientScenarios []clientConnType
 			switch tc.serverTLS {