From 16b6f566793264407b90b6ead49259e3f7731637 Mon Sep 17 00:00:00 2001 From: lhy1024 Date: Tue, 1 Aug 2023 20:06:52 +0800 Subject: [PATCH] add test Signed-off-by: lhy1024 --- pkg/mcs/discovery/register.go | 6 +++++ pkg/mcs/discovery/register_test.go | 38 +++++++++++++++++++++--------- 2 files changed, 33 insertions(+), 11 deletions(-) diff --git a/pkg/mcs/discovery/register.go b/pkg/mcs/discovery/register.go index 35bc9f27f89..958ac62f888 100644 --- a/pkg/mcs/discovery/register.go +++ b/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 } diff --git a/pkg/mcs/discovery/register_test.go b/pkg/mcs/discovery/register_test.go index 6d0bf8a7066..ef6d483b6ff 100644 --- a/pkg/mcs/discovery/register_test.go +++ b/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) }