Fix UT timeout (#22261)

Signed-off-by: Enwei Jiao <enwei.jiao@zilliz.com>

.github/workflows/main.yaml

@@ -72,7 +72,8 @@ jobs:
           echo "corehash=${CORE_HASH}" >> $GITHUB_ENV
           echo "Set CCache hash to ${CORE_HASH}"
       - name: Cache CCache Volumes
-        uses: actions/cache@v3
+        # uses: actions/cache@v3
+        uses: pat-s/always-upload-cache@v3
         with:
           path: .docker/amd64-ubuntu${{ matrix.ubuntu }}-ccache
           key: ubuntu${{ matrix.ubuntu }}-ccache-${{ env.corehash }}
@@ -84,7 +85,8 @@ jobs:
           key: ubuntu${{ matrix.ubuntu }}-go-mod-${{ hashFiles('**/go.sum') }}
           restore-keys: ubuntu${{ matrix.ubuntu }}-go-mod-
       - name: Cache Conan Packages
-        uses: actions/cache@v3
+        # uses: actions/cache@v3
+        uses: pat-s/always-upload-cache@v3
         with:
           path: .docker/amd64-ubuntu${{ matrix.ubuntu }}-conan
           key: ubuntu${{ matrix.ubuntu }}-conan-${{ hashFiles('internal/core/conanfile.*') }}

internal/datacoord/server_test.go

@@ -71,10 +71,10 @@ func TestMain(m *testing.M) {
 	defer embedetcdServer.Close()
 
 	addrs := etcd.GetEmbedEtcdEndpoints(embedetcdServer)
-	// setup env for etcd endpoint
-	os.Setenv("etcd.endpoints", strings.Join(addrs, ","))
 
 	paramtable.Init()
+	paramtable.Get().Save(Params.EtcdCfg.Endpoints.Key, strings.Join(addrs, ","))
+
 	rand.Seed(time.Now().UnixNano())
 	os.Exit(m.Run())
 }

internal/datanode/compaction_executor_test.go

@@ -133,16 +133,21 @@ type mockCompactor struct {
 	isvalid       bool
 	alwaysWorking bool
 
+	mu sync.Mutex
 	wg sync.WaitGroup
 }
 
 var _ compactor = (*mockCompactor)(nil)
 
 func (mc *mockCompactor) start() {
+	mc.mu.Lock()
+	defer mc.mu.Unlock()
 	mc.wg.Add(1)
 }
 
 func (mc *mockCompactor) complete() {
+	mc.mu.Lock()
+	defer mc.mu.Unlock()
 	mc.wg.Done()
 }
 
@@ -164,6 +169,8 @@ func (mc *mockCompactor) getPlanID() UniqueID {
 func (mc *mockCompactor) stop() {
 	if mc.cancel != nil {
 		mc.cancel()
+		mc.mu.Lock()
+		defer mc.mu.Unlock()
 		mc.wg.Wait()
 	}
 }

internal/datanode/data_node_test.go

@@ -69,6 +69,7 @@ func TestMain(t *testing.M) {
 
 	Params.Init()
 	// change to specific channel for test
+	paramtable.Get().Save(Params.EtcdCfg.Endpoints.Key, strings.Join(addrs, ","))
 	paramtable.Get().Save(Params.CommonCfg.DataCoordTimeTick.Key, Params.CommonCfg.DataCoordTimeTick.GetValue()+strconv.Itoa(rand.Int()))
 
 	rateCol, err = newRateCollector()

internal/distributed/connection_manager_test.go

@@ -52,10 +52,9 @@ func TestMain(t *testing.M) {
 	defer embedetcdServer.Server.Stop()
 
 	addrs := etcd.GetEmbedEtcdEndpoints(embedetcdServer)
-	// setup env for etcd endpoint
-	os.Setenv("etcd.endpoints", strings.Join(addrs, ","))
 
 	paramtable.Init()
+	paramtable.Get().Save(paramtable.Get().EtcdCfg.Endpoints.Key, strings.Join(addrs, ","))
 
 	os.Exit(t.Run())
 }

internal/distributed/datacoord/client/client_test.go

@@ -19,16 +19,40 @@ package grpcdatacoordclient
 import (
 	"context"
 	"errors"
+	"math/rand"
+	"os"
+	"strings"
 	"testing"
+	"time"
 
+	"github.com/milvus-io/milvus/internal/log"
 	"github.com/milvus-io/milvus/internal/proto/datapb"
 	"github.com/milvus-io/milvus/internal/proxy"
 	"github.com/milvus-io/milvus/internal/util/etcd"
 	"github.com/milvus-io/milvus/internal/util/mock"
+	"github.com/milvus-io/milvus/internal/util/paramtable"
 	"github.com/stretchr/testify/assert"
 	"google.golang.org/grpc"
 )
 
+func TestMain(m *testing.M) {
+	// init embed etcd
+	embedetcdServer, tempDir, err := etcd.StartTestEmbedEtcdServer()
+	if err != nil {
+		log.Fatal(err.Error())
+	}
+	defer os.RemoveAll(tempDir)
+	defer embedetcdServer.Close()
+
+	addrs := etcd.GetEmbedEtcdEndpoints(embedetcdServer)
+
+	paramtable.Init()
+	paramtable.Get().Save(Params.EtcdCfg.Endpoints.Key, strings.Join(addrs, ","))
+
+	rand.Seed(time.Now().UnixNano())
+	os.Exit(m.Run())
+}
+
 func Test_NewClient(t *testing.T) {
 	proxy.Params.Init()
 

internal/distributed/querycoord/client/client_test.go

@@ -19,10 +19,16 @@ package grpcquerycoordclient
 import (
 	"context"
 	"errors"
+	"math/rand"
+	"os"
+	"strings"
 	"testing"
+	"time"
 
+	"github.com/milvus-io/milvus/internal/log"
 	"github.com/milvus-io/milvus/internal/proto/querypb"
 	"github.com/milvus-io/milvus/internal/util/mock"
+	"github.com/milvus-io/milvus/internal/util/paramtable"
 	"google.golang.org/grpc"
 
 	"github.com/milvus-io/milvus/internal/proxy"
@@ -30,6 +36,24 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+func TestMain(m *testing.M) {
+	// init embed etcd
+	embedetcdServer, tempDir, err := etcd.StartTestEmbedEtcdServer()
+	if err != nil {
+		log.Fatal(err.Error())
+	}
+	defer os.RemoveAll(tempDir)
+	defer embedetcdServer.Close()
+
+	addrs := etcd.GetEmbedEtcdEndpoints(embedetcdServer)
+
+	paramtable.Init()
+	paramtable.Get().Save(Params.EtcdCfg.Endpoints.Key, strings.Join(addrs, ","))
+
+	rand.Seed(time.Now().UnixNano())
+	os.Exit(m.Run())
+}
+
 func Test_NewClient(t *testing.T) {
 	proxy.Params.Init()
 

internal/distributed/rootcoord/client/client_test.go

@@ -19,11 +19,16 @@ package grpcrootcoordclient
 import (
 	"context"
 	"errors"
+	"math/rand"
+	"os"
+	"strings"
 	"testing"
 	"time"
 
+	"github.com/milvus-io/milvus/internal/log"
 	"github.com/milvus-io/milvus/internal/proto/rootcoordpb"
 	"github.com/milvus-io/milvus/internal/util/mock"
+	"github.com/milvus-io/milvus/internal/util/paramtable"
 	"google.golang.org/grpc"
 
 	"github.com/milvus-io/milvus/internal/proxy"
@@ -31,6 +36,24 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+func TestMain(m *testing.M) {
+	// init embed etcd
+	embedetcdServer, tempDir, err := etcd.StartTestEmbedEtcdServer()
+	if err != nil {
+		log.Fatal(err.Error())
+	}
+	defer os.RemoveAll(tempDir)
+	defer embedetcdServer.Close()
+
+	addrs := etcd.GetEmbedEtcdEndpoints(embedetcdServer)
+
+	paramtable.Init()
+	paramtable.Get().Save(Params.EtcdCfg.Endpoints.Key, strings.Join(addrs, ","))
+
+	rand.Seed(time.Now().UnixNano())
+	os.Exit(m.Run())
+}
+
 func Test_NewClient(t *testing.T) {
 	proxy.Params.Init()
 

internal/proxy/impl_test.go

@@ -305,6 +305,7 @@ func TestProxy_InvalidResourceGroupName(t *testing.T) {
 
 	qc := types.NewMockQueryCoord(t)
 	node.SetQueryCoordClient(qc)
+	qc.EXPECT().DropResourceGroup(mock.Anything, mock.Anything).Return(&commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil)
 
 	tsoAllocatorIns := newMockTsoAllocator()
 	node.sched, err = newTaskScheduler(node.ctx, tsoAllocatorIns, node.factory)