From 36279e0797dbd2df4d658c20caceb4c2746e293e Mon Sep 17 00:00:00 2001
From: Danielle Lancashire <dani@builds.terrible.systems>
Date: Wed, 16 Mar 2022 16:00:44 +0000
Subject: [PATCH] tests: migrate TestCtlV3LeaseKeepAliveOnce.* to common

---
 tests/common/lease_test.go     | 52 ++++++++++++++++++++++++++++++++++
 tests/e2e/ctl_v3_lease_test.go | 45 -----------------------------
 2 files changed, 52 insertions(+), 45 deletions(-)

diff --git a/tests/common/lease_test.go b/tests/common/lease_test.go
index caeabf7da..98e382ceb 100644
--- a/tests/common/lease_test.go
+++ b/tests/common/lease_test.go
@@ -207,3 +207,55 @@ func TestLeaseGrantTimeToLiveExpired(t *testing.T) {
 		})
 	}
 }
+
+func TestLeaseGrantKeepAliveOnce(t *testing.T) {
+	testRunner.BeforeTest(t)
+
+	tcs := []struct {
+		name   string
+		config config.ClusterConfig
+	}{
+		{
+			name:   "NoTLS",
+			config: config.ClusterConfig{ClusterSize: 1},
+		},
+		{
+			name:   "PeerTLS",
+			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.ManualTLS},
+		},
+		{
+			name:   "PeerAutoTLS",
+			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.AutoTLS},
+		},
+		{
+			name:   "ClientTLS",
+			config: config.ClusterConfig{ClusterSize: 1, ClientTLS: config.ManualTLS},
+		},
+		{
+			name:   "ClientAutoTLS",
+			config: config.ClusterConfig{ClusterSize: 1, ClientTLS: config.AutoTLS},
+		},
+	}
+	for _, tc := range tcs {
+		t.Run(tc.name, func(t *testing.T) {
+			clus := testRunner.NewCluster(t, tc.config)
+			defer clus.Close()
+			cc := clus.Client()
+
+			testutils.ExecuteWithTimeout(t, 10*time.Second, func() {
+				leaseResp, err := cc.Grant(2)
+				require.NoError(t, err)
+
+				_, err = cc.LeaseKeepAliveOnce(leaseResp.ID)
+				require.NoError(t, err)
+
+				time.Sleep(2 * time.Second) // Wait for the original lease to expire
+
+				ttlResp, err := cc.TimeToLive(leaseResp.ID, config.LeaseOption{})
+				require.NoError(t, err)
+				// We still have a lease!
+				require.Greater(t, int64(2), ttlResp.TTL)
+			})
+		})
+	}
+}
diff --git a/tests/e2e/ctl_v3_lease_test.go b/tests/e2e/ctl_v3_lease_test.go
index 0c1173cd2..905351c93 100644
--- a/tests/e2e/ctl_v3_lease_test.go
+++ b/tests/e2e/ctl_v3_lease_test.go
@@ -37,20 +37,6 @@ func TestCtlV3LeaseKeepAlivePeerTLS(t *testing.T) {
 	testCtl(t, leaseTestKeepAlive, withCfg(*e2e.NewConfigPeerTLS()))
 }
 
-func TestCtlV3LeaseKeepAliveOnce(t *testing.T) { testCtl(t, leaseTestKeepAliveOnce) }
-func TestCtlV3LeaseKeepAliveOnceNoTLS(t *testing.T) {
-	testCtl(t, leaseTestKeepAliveOnce, withCfg(*e2e.NewConfigNoTLS()))
-}
-func TestCtlV3LeaseKeepAliveOnceClientTLS(t *testing.T) {
-	testCtl(t, leaseTestKeepAliveOnce, withCfg(*e2e.NewConfigClientTLS()))
-}
-func TestCtlV3LeaseKeepAliveOnceClientAutoTLS(t *testing.T) {
-	testCtl(t, leaseTestKeepAliveOnce, withCfg(*e2e.NewConfigClientAutoTLS()))
-}
-func TestCtlV3LeaseKeepAliveOncePeerTLS(t *testing.T) {
-	testCtl(t, leaseTestKeepAliveOnce, withCfg(*e2e.NewConfigPeerTLS()))
-}
-
 func TestCtlV3LeaseRevoke(t *testing.T) { testCtl(t, leaseTestRevoked) }
 func TestCtlV3LeaseRevokeNoTLS(t *testing.T) {
 	testCtl(t, leaseTestRevoked, withCfg(*e2e.NewConfigNoTLS()))
@@ -82,23 +68,6 @@ func leaseTestKeepAlive(cx ctlCtx) {
 	}
 }
 
-func leaseTestKeepAliveOnce(cx ctlCtx) {
-	// put with TTL 10 seconds and keep-alive once
-	leaseID, err := ctlV3LeaseGrant(cx, 10)
-	if err != nil {
-		cx.t.Fatalf("leaseTestKeepAlive: ctlV3LeaseGrant error (%v)", err)
-	}
-	if err := ctlV3Put(cx, "key", "val", leaseID); err != nil {
-		cx.t.Fatalf("leaseTestKeepAlive: ctlV3Put error (%v)", err)
-	}
-	if err := ctlV3LeaseKeepAliveOnce(cx, leaseID); err != nil {
-		cx.t.Fatalf("leaseTestKeepAlive: ctlV3LeaseKeepAliveOnce error (%v)", err)
-	}
-	if err := ctlV3Get(cx, []string{"key"}, kv{"key", "val"}); err != nil {
-		cx.t.Fatalf("leaseTestKeepAlive: ctlV3Get error (%v)", err)
-	}
-}
-
 func leaseTestRevoked(cx ctlCtx) {
 	err := leaseTestRevoke(cx)
 	if err != nil {
@@ -161,20 +130,6 @@ func ctlV3LeaseKeepAlive(cx ctlCtx, leaseID string) error {
 	return proc.Stop()
 }
 
-func ctlV3LeaseKeepAliveOnce(cx ctlCtx, leaseID string) error {
-	cmdArgs := append(cx.PrefixArgs(), "lease", "keep-alive", "--once", leaseID)
-
-	proc, err := e2e.SpawnCmd(cmdArgs, nil)
-	if err != nil {
-		return err
-	}
-
-	if _, err = proc.Expect(fmt.Sprintf("lease %s keepalived with TTL(", leaseID)); err != nil {
-		return err
-	}
-	return proc.Stop()
-}
-
 func ctlV3LeaseRevoke(cx ctlCtx, leaseID string) error {
 	cmdArgs := append(cx.PrefixArgs(), "lease", "revoke", leaseID)
 	return e2e.SpawnWithExpectWithEnv(cmdArgs, cx.envMap, fmt.Sprintf("lease %s revoked", leaseID))