Use embed etcd for ConnectManager unit test (#21884) (#21902)

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>

internal/distributed/connection_manager_test.go

@@ -42,8 +42,26 @@ import (
 
 var Params paramtable.BaseTable
 
+func TestMain(t *testing.M) {
+	// init embed etcd
+	embedetcdServer, tempDir, err := etcd.StartTestEmbedEtcdServer()
+	if err != nil {
+		log.Fatal("failed to start embed etcd server for unittest", zap.Error(err))
+	}
+
+	defer os.RemoveAll(tempDir)
+	defer embedetcdServer.Server.Stop()
+
+	addrs := etcd.GetEmbedEtcdEndpoints(embedetcdServer)
+	// setup env for etcd endpoint
+	os.Setenv("etcd.endpoints", strings.Join(addrs, ","))
+
+	os.Exit(t.Run())
+}
+
 func TestConnectionManager(t *testing.T) {
 	Params.Init()
+
 	ctx := context.Background()
 
 	session := initSession(ctx)

internal/util/etcd/etcd_util.go

@@ -21,11 +21,14 @@ import (
 	"crypto/x509"
 	"fmt"
 	"io/ioutil"
+	"net/url"
+	"os"
 	"time"
 
 	"github.com/pkg/errors"
 
 	clientv3 "go.etcd.io/etcd/client/v3"
+	"go.etcd.io/etcd/server/v3/embed"
 )
 
 var (
@@ -164,3 +167,39 @@ func buildKvGroup(keys, values []string) (map[string]string, error) {
 	}
 	return ret, nil
 }
+
+// StartTestEmbedEtcdServer returns a newly created embed etcd server.
+// ### USED FOR UNIT TEST ONLY ###
+func StartTestEmbedEtcdServer() (*embed.Etcd, string, error) {
+	dir, err := ioutil.TempDir(os.TempDir(), "milvus_datanode_ut")
+	if err != nil {
+		return nil, "", err
+	}
+	config := embed.NewConfig()
+
+	config.Dir = dir
+	config.LogLevel = "warn"
+	config.LogOutputs = []string{"default"}
+	u, err := url.Parse("http://localhost:0")
+	if err != nil {
+		return nil, "", err
+	}
+	config.LCUrls = []url.URL{*u}
+	u, err = url.Parse("http://localhost:0")
+	if err != nil {
+		return nil, "", err
+	}
+	config.LPUrls = []url.URL{*u}
+
+	server, err := embed.StartEtcd(config)
+	return server, dir, err
+}
+
+// GetEmbedEtcdEndpoints returns etcd listener address for endpoint config.
+func GetEmbedEtcdEndpoints(server *embed.Etcd) []string {
+	addrs := make([]string, 0, len(server.Clients))
+	for _, l := range server.Clients {
+		addrs = append(addrs, l.Addr().String())
+	}
+	return addrs
+}