From 12581188e831d3ec866bdd8fe00e4409338a25ef Mon Sep 17 00:00:00 2001
From: Derek Nola <derek.nola@suse.com>
Date: Tue, 30 Nov 2021 11:05:13 -0800
Subject: [PATCH] Improved cleanup for etcd unit test (#4537) (#4608)

* Improved cleanup for etcd unit test

Signed-off-by: Derek Nola <derek.nola@suse.com>
---
 pkg/etcd/etcd_test.go | 73 +++++++++++++++++++++++++++++++------------
 1 file changed, 53 insertions(+), 20 deletions(-)

diff --git a/pkg/etcd/etcd_test.go b/pkg/etcd/etcd_test.go
index d7f7cd4153..8fce94b3e3 100644
--- a/pkg/etcd/etcd_test.go
+++ b/pkg/etcd/etcd_test.go
@@ -14,6 +14,7 @@ import (
 	testutil "github.com/rancher/k3s/tests/util"
 	"github.com/robfig/cron/v3"
 	clientv3 "go.etcd.io/etcd/client/v3"
+	"go.etcd.io/etcd/server/v3/etcdserver"
 	utilnet "k8s.io/apimachinery/pkg/util/net"
 )
 
@@ -222,8 +223,8 @@ func Test_UnitETCD_Start(t *testing.T) {
 		name     string
 		fields   fields
 		args     args
-		setup    func(cnf *config.Control, ctxInfo *contextInfo) error
-		teardown func(cnf *config.Control, ctxInfo *contextInfo) error
+		setup    func(e *ETCD, ctxInfo *contextInfo) error
+		teardown func(e *ETCD, ctxInfo *contextInfo) error
 		wantErr  bool
 	}{
 		{
@@ -236,15 +237,24 @@ func Test_UnitETCD_Start(t *testing.T) {
 			args: args{
 				clientAccessInfo: nil,
 			},
-			setup: func(cnf *config.Control, ctxInfo *contextInfo) error {
+			setup: func(e *ETCD, ctxInfo *contextInfo) error {
 				ctxInfo.ctx, ctxInfo.cancel = context.WithCancel(context.Background())
-				cnf.EtcdDisableSnapshots = true
-				return testutil.GenerateRuntime(cnf)
+				e.config.EtcdDisableSnapshots = true
+				testutil.GenerateRuntime(e.config)
+				e.runtime = e.config.Runtime
+				client, err := GetClient(ctxInfo.ctx, e.runtime, endpoint)
+				e.client = client
+
+				return err
 			},
-			teardown: func(cnf *config.Control, ctxInfo *contextInfo) error {
+			teardown: func(e *ETCD, ctxInfo *contextInfo) error {
+				// RemoveSelf will fail with a specific error, but it still does cleanup for testing purposes
+				if err := e.RemoveSelf(ctxInfo.ctx); err != nil && err.Error() != etcdserver.ErrNotEnoughStartedMembers.Error() {
+					return err
+				}
 				ctxInfo.cancel()
-				time.Sleep(5 * time.Second)
-				testutil.CleanupDataDir(cnf)
+				time.Sleep(10 * time.Second)
+				testutil.CleanupDataDir(e.config)
 				return nil
 			},
 		},
@@ -259,14 +269,23 @@ func Test_UnitETCD_Start(t *testing.T) {
 			args: args{
 				clientAccessInfo: nil,
 			},
-			setup: func(cnf *config.Control, ctxInfo *contextInfo) error {
+			setup: func(e *ETCD, ctxInfo *contextInfo) error {
 				ctxInfo.ctx, ctxInfo.cancel = context.WithCancel(context.Background())
-				return testutil.GenerateRuntime(cnf)
+				testutil.GenerateRuntime(e.config)
+				e.runtime = e.config.Runtime
+				client, err := GetClient(ctxInfo.ctx, e.runtime, endpoint)
+				e.client = client
+
+				return err
 			},
-			teardown: func(cnf *config.Control, ctxInfo *contextInfo) error {
+			teardown: func(e *ETCD, ctxInfo *contextInfo) error {
+				// RemoveSelf will fail with a specific error, but it still does cleanup for testing purposes
+				if err := e.RemoveSelf(ctxInfo.ctx); err != nil && err.Error() != etcdserver.ErrNotEnoughStartedMembers.Error() {
+					return err
+				}
 				ctxInfo.cancel()
 				time.Sleep(5 * time.Second)
-				testutil.CleanupDataDir(cnf)
+				testutil.CleanupDataDir(e.config)
 				return nil
 			},
 		},
@@ -281,18 +300,28 @@ func Test_UnitETCD_Start(t *testing.T) {
 			args: args{
 				clientAccessInfo: nil,
 			},
-			setup: func(cnf *config.Control, ctxInfo *contextInfo) error {
+			setup: func(e *ETCD, ctxInfo *contextInfo) error {
 				ctxInfo.ctx, ctxInfo.cancel = context.WithCancel(context.Background())
-				if err := testutil.GenerateRuntime(cnf); err != nil {
+				if err := testutil.GenerateRuntime(e.config); err != nil {
 					return err
 				}
-				return os.MkdirAll(walDir(cnf), 0700)
+				e.runtime = e.config.Runtime
+				client, err := GetClient(ctxInfo.ctx, e.runtime, endpoint)
+				if err != nil {
+					return err
+				}
+				e.client = client
+				return os.MkdirAll(walDir(e.config), 0700)
 			},
-			teardown: func(cnf *config.Control, ctxInfo *contextInfo) error {
+			teardown: func(e *ETCD, ctxInfo *contextInfo) error {
+				// RemoveSelf will fail with a specific error, but it still does cleanup for testing purposes
+				if err := e.RemoveSelf(ctxInfo.ctx); err != nil && err.Error() != etcdserver.ErrNotEnoughStartedMembers.Error() {
+					return err
+				}
 				ctxInfo.cancel()
 				time.Sleep(5 * time.Second)
-				testutil.CleanupDataDir(cnf)
-				os.Remove(walDir(cnf))
+				testutil.CleanupDataDir(e.config)
+				os.Remove(walDir(e.config))
 				return nil
 			},
 		},
@@ -308,14 +337,18 @@ func Test_UnitETCD_Start(t *testing.T) {
 				cron:    tt.fields.cron,
 				s3:      tt.fields.s3,
 			}
-			defer tt.teardown(e.config, &tt.fields.context)
-			if err := tt.setup(e.config, &tt.fields.context); err != nil {
+
+			if err := tt.setup(e, &tt.fields.context); err != nil {
 				t.Errorf("Setup for ETCD.Start() failed = %v", err)
 				return
 			}
 			if err := e.Start(tt.fields.context.ctx, tt.args.clientAccessInfo); (err != nil) != tt.wantErr {
 				t.Errorf("ETCD.Start() error = %v, wantErr %v", err, tt.wantErr)
 			}
+			if err := tt.teardown(e, &tt.fields.context); err != nil {
+				t.Errorf("Teardown for ETCD.Start() failed = %v", err)
+				return
+			}
 		})
 	}
 }