add test

Signed-off-by: lhy1024 <[email protected]>

pkg/mcs/discovery/register.go

@@ -102,6 +102,12 @@ func (sr *ServiceRegister) Register() error {
 							log.Error("put the key failed", zap.String("key", sr.key), zap.Error(err))
 							continue
 						}
+
+						kresp, err = sr.cli.KeepAlive(sr.ctx, resp.ID)
+						if err != nil {
+							log.Error("client keep alive failed", zap.String("key", sr.key), zap.Error(err))
+							continue
+						}
 						t.Stop()
 						goto KeepaliveLoop
 					}

pkg/mcs/discovery/register_test.go

@@ -29,17 +29,13 @@ func TestRegister(t *testing.T) {
 	re := require.New(t)
 	cfg := etcdutil.NewTestSingleConfig(t)
 	etcd, err := embed.StartEtcd(cfg)
-	defer func() {
-		etcd.Close()
-	}()
 	re.NoError(err)
-
 	ep := cfg.LCUrls[0].String()
 	client, err := clientv3.NewFromURL(ep)
 	re.NoError(err)
-
 	<-etcd.Server.ReadyNotify()
-	// with http prefix
+
+	// Test register with http prefix.
 	sr := NewServiceRegister(context.Background(), client, "12345", "test_service", "http://127.0.0.1:1", "http://127.0.0.1:1", 10)
 	re.NoError(err)
 	err = sr.Register()
@@ -49,20 +45,40 @@ func TestRegister(t *testing.T) {
 	re.NoError(err)
 	re.Equal("http://127.0.0.1:1", string(resp.Kvs[0].Value))
 
+	// Test deregister.
 	err = sr.Deregister()
 	re.NoError(err)
 	resp, err = client.Get(context.Background(), sr.key)
 	re.NoError(err)
 	re.Empty(resp.Kvs)
 
+	// Test the case that ctx is canceled.
 	sr = NewServiceRegister(context.Background(), client, "12345", "test_service", "127.0.0.1:2", "127.0.0.1:2", 1)
-	re.NoError(err)
 	err = sr.Register()
 	re.NoError(err)
 	sr.cancel()
-	// ensure that the lease is expired
-	time.Sleep(3 * time.Second)
-	resp, err = client.Get(context.Background(), sr.key)
+	re.Empty(getKeyAfterLeaseExpired(re, client, sr.key))
+
+	// Test the case that keepalive is failed when the etcd is restarted.
+	sr = NewServiceRegister(context.Background(), client, "12345", "test_service", "127.0.0.1:2", "127.0.0.1:2", 1)
+	err = sr.Register()
 	re.NoError(err)
-	re.Empty(resp.Kvs)
+	re.Equal("127.0.0.1:2", getKeyAfterLeaseExpired(re, client, sr.key))
+	etcd.Close()                // close the etcd to make the keepalive failed
+	time.Sleep(3 * time.Second) // ensure that the lease is expired
+	etcd, err = embed.StartEtcd(cfg)
+	re.NoError(err)
+	<-etcd.Server.ReadyNotify()
+	defer etcd.Close()
+	re.Equal("127.0.0.1:2", getKeyAfterLeaseExpired(re, client, sr.key))
+}
+
+func getKeyAfterLeaseExpired(re *require.Assertions, client *clientv3.Client, key string) string {
+	time.Sleep(3 * time.Second) // ensure that the lease is expired
+	resp, err := client.Get(context.Background(), key)
+	re.NoError(err)
+	if len(resp.Kvs) == 0 {
+		return ""
+	}
+	return string(resp.Kvs[0].Value)
 }