From b7a97fe8b76bdf90cc1c5bfa84d0adf40dc3da79 Mon Sep 17 00:00:00 2001
From: congqixia <congqi.xia@zilliz.com>
Date: Wed, 1 Feb 2023 15:59:51 +0800
Subject: [PATCH] Use embed etcd for ConnectManager unit test (#21884) (#21902)

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
---
 .../distributed/connection_manager_test.go    | 18 +++++++++
 internal/util/etcd/etcd_util.go               | 39 +++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/internal/distributed/connection_manager_test.go b/internal/distributed/connection_manager_test.go
index e773e7c5b4..e5c43f5369 100644
--- a/internal/distributed/connection_manager_test.go
+++ b/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)
diff --git a/internal/util/etcd/etcd_util.go b/internal/util/etcd/etcd_util.go
index 60e0488596..2b8f76e61c 100644
--- a/internal/util/etcd/etcd_util.go
+++ b/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
+}