// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
//	http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package meta

import (
	"context"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/suite"

	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
	"github.com/milvus-io/milvus-proto/go-api/v2/rgpb"
	"github.com/milvus-io/milvus/internal/json"
	etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
	"github.com/milvus-io/milvus/internal/kv/mocks"
	"github.com/milvus-io/milvus/internal/metastore/kv/querycoord"
	"github.com/milvus-io/milvus/internal/querycoordv2/params"
	"github.com/milvus-io/milvus/internal/querycoordv2/session"
	"github.com/milvus-io/milvus/internal/util/sessionutil"
	"github.com/milvus-io/milvus/pkg/v2/kv"
	"github.com/milvus-io/milvus/pkg/v2/log"
	"github.com/milvus-io/milvus/pkg/v2/util/etcd"
	"github.com/milvus-io/milvus/pkg/v2/util/merr"
	"github.com/milvus-io/milvus/pkg/v2/util/metricsinfo"
	"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
	"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)

type ResourceManagerSuite struct {
	suite.Suite

	kv      kv.MetaKv
	manager *ResourceManager
	ctx     context.Context
}

func (suite *ResourceManagerSuite) SetupSuite() {
	paramtable.Init()
}

func (suite *ResourceManagerSuite) SetupTest() {
	config := params.GenerateEtcdConfig()
	cli, err := etcd.GetEtcdClient(
		config.UseEmbedEtcd.GetAsBool(),
		config.EtcdUseSSL.GetAsBool(),
		config.Endpoints.GetAsStrings(),
		config.EtcdTLSCert.GetValue(),
		config.EtcdTLSKey.GetValue(),
		config.EtcdTLSCACert.GetValue(),
		config.EtcdTLSMinVersion.GetValue())
	suite.Require().NoError(err)
	suite.kv = etcdkv.NewEtcdKV(cli, config.MetaRootPath.GetValue())

	store := querycoord.NewCatalog(suite.kv)
	suite.manager = NewResourceManager(store, session.NewNodeManager())
	suite.ctx = context.Background()
}

func (suite *ResourceManagerSuite) TearDownSuite() {
	suite.kv.Close()
}

func TestResourceManager(t *testing.T) {
	suite.Run(t, new(ResourceManagerSuite))
}

func (suite *ResourceManagerSuite) TestValidateConfiguration() {
	ctx := suite.ctx
	err := suite.manager.validateResourceGroupConfig("rg1", newResourceGroupConfig(0, 0))
	suite.NoError(err)

	err = suite.manager.validateResourceGroupConfig("rg1", &rgpb.ResourceGroupConfig{})
	suite.ErrorIs(err, merr.ErrResourceGroupIllegalConfig)

	err = suite.manager.validateResourceGroupConfig("rg1", newResourceGroupConfig(-1, 2))
	suite.ErrorIs(err, merr.ErrResourceGroupIllegalConfig)

	err = suite.manager.validateResourceGroupConfig("rg1", newResourceGroupConfig(2, -1))
	suite.ErrorIs(err, merr.ErrResourceGroupIllegalConfig)

	err = suite.manager.validateResourceGroupConfig("rg1", newResourceGroupConfig(3, 2))
	suite.ErrorIs(err, merr.ErrResourceGroupIllegalConfig)

	cfg := newResourceGroupConfig(0, 0)
	cfg.TransferFrom = []*rgpb.ResourceGroupTransfer{{ResourceGroup: "rg1"}}
	err = suite.manager.validateResourceGroupConfig("rg1", cfg)
	suite.ErrorIs(err, merr.ErrResourceGroupIllegalConfig)

	cfg = newResourceGroupConfig(0, 0)
	cfg.TransferFrom = []*rgpb.ResourceGroupTransfer{{ResourceGroup: "rg2"}}
	err = suite.manager.validateResourceGroupConfig("rg1", cfg)
	suite.ErrorIs(err, merr.ErrResourceGroupIllegalConfig)

	cfg = newResourceGroupConfig(0, 0)
	cfg.TransferTo = []*rgpb.ResourceGroupTransfer{{ResourceGroup: "rg1"}}
	err = suite.manager.validateResourceGroupConfig("rg1", cfg)
	suite.ErrorIs(err, merr.ErrResourceGroupIllegalConfig)

	cfg = newResourceGroupConfig(0, 0)
	cfg.TransferTo = []*rgpb.ResourceGroupTransfer{{ResourceGroup: "rg2"}}
	err = suite.manager.validateResourceGroupConfig("rg1", cfg)
	suite.ErrorIs(err, merr.ErrResourceGroupIllegalConfig)

	err = suite.manager.AddResourceGroup(ctx, "rg2", newResourceGroupConfig(0, 0))
	suite.NoError(err)

	err = suite.manager.RemoveResourceGroup(ctx, "rg2")
	suite.NoError(err)
}

func (suite *ResourceManagerSuite) TestValidateDelete() {
	ctx := suite.ctx
	// Non empty resource group can not be removed.
	err := suite.manager.AddResourceGroup(ctx, "rg1", newResourceGroupConfig(1, 1))
	suite.NoError(err)

	err = suite.manager.validateResourceGroupIsDeletable(DefaultResourceGroupName)
	suite.ErrorIs(err, merr.ErrParameterInvalid)

	err = suite.manager.validateResourceGroupIsDeletable("rg1")
	suite.ErrorIs(err, merr.ErrParameterInvalid)

	cfg := newResourceGroupConfig(0, 0)
	cfg.TransferFrom = []*rgpb.ResourceGroupTransfer{{ResourceGroup: "rg1"}}
	suite.manager.AddResourceGroup(ctx, "rg2", cfg)
	suite.manager.UpdateResourceGroups(ctx, map[string]*rgpb.ResourceGroupConfig{
		"rg1": newResourceGroupConfig(0, 0),
	})
	err = suite.manager.validateResourceGroupIsDeletable("rg1")
	suite.ErrorIs(err, merr.ErrParameterInvalid)

	cfg = newResourceGroupConfig(0, 0)
	cfg.TransferTo = []*rgpb.ResourceGroupTransfer{{ResourceGroup: "rg1"}}
	suite.manager.UpdateResourceGroups(ctx, map[string]*rgpb.ResourceGroupConfig{
		"rg2": cfg,
	})
	err = suite.manager.validateResourceGroupIsDeletable("rg1")
	suite.ErrorIs(err, merr.ErrParameterInvalid)

	suite.manager.UpdateResourceGroups(ctx, map[string]*rgpb.ResourceGroupConfig{
		"rg2": newResourceGroupConfig(0, 0),
	})
	err = suite.manager.validateResourceGroupIsDeletable("rg1")
	suite.NoError(err)

	err = suite.manager.RemoveResourceGroup(ctx, "rg1")
	suite.NoError(err)
	err = suite.manager.RemoveResourceGroup(ctx, "rg2")
	suite.NoError(err)
}

func (suite *ResourceManagerSuite) TestManipulateResourceGroup() {
	ctx := suite.ctx
	// test add rg
	err := suite.manager.AddResourceGroup(ctx, "rg1", newResourceGroupConfig(0, 0))
	suite.NoError(err)
	suite.True(suite.manager.ContainResourceGroup(ctx, "rg1"))
	suite.Len(suite.manager.ListResourceGroups(ctx), 2)

	// test add duplicate rg but same configuration is ok
	err = suite.manager.AddResourceGroup(ctx, "rg1", newResourceGroupConfig(0, 0))
	suite.NoError(err)

	err = suite.manager.AddResourceGroup(ctx, "rg1", newResourceGroupConfig(1, 1))
	suite.Error(err)

	// test delete rg
	err = suite.manager.RemoveResourceGroup(ctx, "rg1")
	suite.NoError(err)

	// test delete rg which doesn't exist
	err = suite.manager.RemoveResourceGroup(ctx, "rg1")
	suite.NoError(err)
	// test delete default rg
	err = suite.manager.RemoveResourceGroup(ctx, DefaultResourceGroupName)
	suite.ErrorIs(err, merr.ErrParameterInvalid)

	// test delete a rg not empty.
	err = suite.manager.AddResourceGroup(ctx, "rg2", newResourceGroupConfig(1, 1))
	suite.NoError(err)
	err = suite.manager.RemoveResourceGroup(ctx, "rg2")
	suite.ErrorIs(err, merr.ErrParameterInvalid)

	// test delete a rg after update
	suite.manager.UpdateResourceGroups(ctx, map[string]*rgpb.ResourceGroupConfig{
		"rg2": newResourceGroupConfig(0, 0),
	})
	err = suite.manager.RemoveResourceGroup(ctx, "rg2")
	suite.NoError(err)

	// assign a node to rg.
	err = suite.manager.AddResourceGroup(ctx, "rg2", newResourceGroupConfig(1, 1))
	suite.NoError(err)
	suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   1,
		Address:  "localhost",
		Hostname: "localhost",
	}))
	defer suite.manager.nodeMgr.Remove(1)
	suite.manager.HandleNodeUp(ctx, 1)
	err = suite.manager.RemoveResourceGroup(ctx, "rg2")
	suite.ErrorIs(err, merr.ErrParameterInvalid)
	suite.manager.UpdateResourceGroups(ctx, map[string]*rgpb.ResourceGroupConfig{
		"rg2": newResourceGroupConfig(0, 0),
	})
	log.Info("xxxxx")
	// RemoveResourceGroup will remove all nodes from the resource group.
	err = suite.manager.RemoveResourceGroup(ctx, "rg2")
	suite.NoError(err)

	suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   10,
		Address:  "localhost",
		Hostname: "localhost",
		Labels: map[string]string{
			sessionutil.LabelStreamingNodeEmbeddedQueryNode: "1",
		},
	}))
	suite.manager.HandleNodeUp(ctx, 10)
}

func (suite *ResourceManagerSuite) TestNodeUpAndDown() {
	ctx := suite.ctx
	suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   1,
		Address:  "localhost",
		Hostname: "localhost",
	}))
	err := suite.manager.AddResourceGroup(ctx, "rg1", newResourceGroupConfig(1, 1))
	suite.NoError(err)
	// test add node to rg
	suite.manager.HandleNodeUp(ctx, 1)
	suite.Equal(1, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())

	// test add non-exist node to rg
	err = suite.manager.UpdateResourceGroups(ctx, map[string]*rgpb.ResourceGroupConfig{
		"rg1": newResourceGroupConfig(2, 3),
	})
	suite.NoError(err)
	suite.manager.HandleNodeUp(ctx, 2)
	suite.Equal(1, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Zero(suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())

	// teardown a non-exist node from rg.
	suite.manager.HandleNodeDown(ctx, 2)
	suite.Equal(1, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Zero(suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())

	// test add exist node to rg
	suite.manager.HandleNodeUp(ctx, 1)
	suite.Equal(1, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Zero(suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())

	// teardown a exist node from rg.
	suite.manager.HandleNodeDown(ctx, 1)
	suite.Zero(suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Zero(suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())

	// teardown a exist node from rg.
	suite.manager.HandleNodeDown(ctx, 1)
	suite.Zero(suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Zero(suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())

	suite.manager.HandleNodeUp(ctx, 1)
	suite.Equal(1, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Zero(suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())

	err = suite.manager.UpdateResourceGroups(ctx, map[string]*rgpb.ResourceGroupConfig{
		"rg1": newResourceGroupConfig(4, 4),
	})
	suite.NoError(err)
	suite.manager.AddResourceGroup(ctx, "rg2", newResourceGroupConfig(1, 1))
	suite.NoError(err)

	suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   11,
		Address:  "localhost",
		Hostname: "localhost",
	}))
	suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   12,
		Address:  "localhost",
		Hostname: "localhost",
	}))
	suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   13,
		Address:  "localhost",
		Hostname: "localhost",
	}))
	suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   14,
		Address:  "localhost",
		Hostname: "localhost",
	}))
	suite.manager.HandleNodeUp(ctx, 11)
	suite.manager.HandleNodeUp(ctx, 12)
	suite.manager.HandleNodeUp(ctx, 13)
	suite.manager.HandleNodeUp(ctx, 14)

	suite.Equal(4, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Equal(1, suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	suite.Zero(suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())

	suite.manager.HandleNodeDown(ctx, 11)
	suite.manager.HandleNodeDown(ctx, 12)
	suite.manager.HandleNodeDown(ctx, 13)
	suite.manager.HandleNodeDown(ctx, 14)
	suite.Equal(1, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Zero(suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	suite.Zero(suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())

	suite.manager.HandleNodeDown(ctx, 1)
	suite.Zero(suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Zero(suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	suite.Zero(suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())

	suite.manager.UpdateResourceGroups(ctx, map[string]*rgpb.ResourceGroupConfig{
		"rg1": newResourceGroupConfig(20, 30),
		"rg2": newResourceGroupConfig(30, 40),
	})
	for i := 1; i <= 100; i++ {
		suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
			NodeID:   int64(i),
			Address:  "localhost",
			Hostname: "localhost",
		}))
		suite.manager.HandleNodeUp(ctx, int64(i))
	}

	suite.Equal(20, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Equal(30, suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	suite.Equal(50, suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())

	// down all nodes
	for i := 1; i <= 100; i++ {
		suite.manager.HandleNodeDown(ctx, int64(i))
		suite.Equal(100-i, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum()+
			suite.manager.GetResourceGroup(ctx, "rg2").NodeNum()+
			suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())
	}

	// if there are all rgs reach limit, should be fall back to default rg.
	suite.manager.UpdateResourceGroups(ctx, map[string]*rgpb.ResourceGroupConfig{
		"rg1":                    newResourceGroupConfig(0, 0),
		"rg2":                    newResourceGroupConfig(0, 0),
		DefaultResourceGroupName: newResourceGroupConfig(0, 0),
	})

	for i := 1; i <= 100; i++ {
		suite.manager.HandleNodeUp(ctx, int64(i))
		suite.Equal(i, suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())
		suite.Equal(0, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
		suite.Equal(0, suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	}
}

func (suite *ResourceManagerSuite) TestAutoRecover() {
	ctx := suite.ctx
	for i := 1; i <= 100; i++ {
		suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
			NodeID:   int64(i),
			Address:  "localhost",
			Hostname: "localhost",
		}))
		suite.manager.HandleNodeUp(ctx, int64(i))
	}
	suite.Equal(100, suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())

	// Recover 10 nodes from default resource group
	suite.manager.AddResourceGroup(ctx, "rg1", newResourceGroupConfig(10, 30))
	suite.Zero(suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg1").MissingNumOfNodes())
	suite.Equal(100, suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())
	suite.manager.AutoRecoverResourceGroup(ctx, "rg1")
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Equal(0, suite.manager.GetResourceGroup(ctx, "rg1").MissingNumOfNodes())
	suite.Equal(90, suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())

	// Recover 20 nodes from default resource group
	suite.manager.AddResourceGroup(ctx, "rg2", newResourceGroupConfig(20, 30))
	suite.Zero(suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	suite.Equal(20, suite.manager.GetResourceGroup(ctx, "rg2").MissingNumOfNodes())
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Equal(90, suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())
	suite.manager.AutoRecoverResourceGroup(ctx, "rg2")
	suite.Equal(20, suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Equal(70, suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())

	// Recover 5 redundant nodes from resource group
	suite.manager.UpdateResourceGroups(ctx, map[string]*rgpb.ResourceGroupConfig{
		"rg1": newResourceGroupConfig(5, 5),
	})
	suite.manager.AutoRecoverResourceGroup(ctx, "rg1")
	suite.Equal(20, suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	suite.Equal(5, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Equal(75, suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())

	// Recover 10 redundant nodes from resource group 2 to resource group 1 and default resource group.
	suite.manager.UpdateResourceGroups(ctx, map[string]*rgpb.ResourceGroupConfig{
		"rg1": newResourceGroupConfig(10, 20),
		"rg2": newResourceGroupConfig(5, 10),
	})

	suite.manager.AutoRecoverResourceGroup(ctx, "rg2")
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	suite.Equal(80, suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())

	// recover redundant nodes from default resource group
	suite.manager.UpdateResourceGroups(ctx, map[string]*rgpb.ResourceGroupConfig{
		"rg1":                    newResourceGroupConfig(10, 20),
		"rg2":                    newResourceGroupConfig(20, 30),
		DefaultResourceGroupName: newResourceGroupConfig(10, 20),
	})
	suite.manager.AutoRecoverResourceGroup(ctx, "rg1")
	suite.manager.AutoRecoverResourceGroup(ctx, "rg2")
	suite.manager.AutoRecoverResourceGroup(ctx, DefaultResourceGroupName)

	// Even though the default resource group has 20 nodes limits,
	// all redundant nodes will be assign to default resource group.
	suite.Equal(20, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Equal(30, suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	suite.Equal(50, suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())

	// Test recover missing from high priority resource group by set `from`.
	suite.manager.AddResourceGroup(ctx, "rg3", &rgpb.ResourceGroupConfig{
		Requests: &rgpb.ResourceGroupLimit{
			NodeNum: 15,
		},
		Limits: &rgpb.ResourceGroupLimit{
			NodeNum: 15,
		},
		TransferFrom: []*rgpb.ResourceGroupTransfer{{
			ResourceGroup: "rg1",
		}},
	})
	suite.manager.UpdateResourceGroups(ctx, map[string]*rgpb.ResourceGroupConfig{
		DefaultResourceGroupName: newResourceGroupConfig(30, 40),
	})

	suite.manager.AutoRecoverResourceGroup(ctx, "rg1")
	suite.manager.AutoRecoverResourceGroup(ctx, "rg2")
	suite.manager.AutoRecoverResourceGroup(ctx, DefaultResourceGroupName)
	suite.manager.AutoRecoverResourceGroup(ctx, "rg3")

	// Get 10 from default group for redundant nodes, get 5 from rg1 for rg3 at high priority.
	suite.Equal(15, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Equal(30, suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	suite.Equal(15, suite.manager.GetResourceGroup(ctx, "rg3").NodeNum())
	suite.Equal(40, suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())

	// Test recover redundant to high priority resource group by set `to`.
	suite.manager.UpdateResourceGroups(ctx, map[string]*rgpb.ResourceGroupConfig{
		"rg3": {
			Requests: &rgpb.ResourceGroupLimit{
				NodeNum: 0,
			},
			Limits: &rgpb.ResourceGroupLimit{
				NodeNum: 0,
			},
			TransferTo: []*rgpb.ResourceGroupTransfer{{
				ResourceGroup: "rg2",
			}},
		},
		"rg1": newResourceGroupConfig(15, 100),
		"rg2": newResourceGroupConfig(15, 40),
	})

	suite.manager.AutoRecoverResourceGroup(ctx, "rg1")
	suite.manager.AutoRecoverResourceGroup(ctx, "rg2")
	suite.manager.AutoRecoverResourceGroup(ctx, DefaultResourceGroupName)
	suite.manager.AutoRecoverResourceGroup(ctx, "rg3")

	// Recover rg3 by transfer 10 nodes to rg2 with high priority, 5 to rg1.
	suite.Equal(20, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Equal(40, suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	suite.Equal(0, suite.manager.GetResourceGroup(ctx, "rg3").NodeNum())
	suite.Equal(40, suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())

	suite.testTransferNode()

	// Test redundant nodes recover to default resource group.
	suite.manager.UpdateResourceGroups(ctx, map[string]*rgpb.ResourceGroupConfig{
		DefaultResourceGroupName: newResourceGroupConfig(1, 1),
		"rg3":                    newResourceGroupConfig(0, 0),
		"rg2":                    newResourceGroupConfig(0, 0),
		"rg1":                    newResourceGroupConfig(0, 0),
	})
	// Even default resource group has 1 node limit,
	// all redundant nodes will be assign to default resource group if there's no resource group can hold.
	suite.manager.AutoRecoverResourceGroup(ctx, DefaultResourceGroupName)
	suite.manager.AutoRecoverResourceGroup(ctx, "rg1")
	suite.manager.AutoRecoverResourceGroup(ctx, "rg2")
	suite.manager.AutoRecoverResourceGroup(ctx, "rg3")
	suite.Equal(0, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Equal(0, suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	suite.Equal(0, suite.manager.GetResourceGroup(ctx, "rg3").NodeNum())
	suite.Equal(100, suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())

	// Test redundant recover to missing nodes and missing nodes from redundant nodes.
	// Initialize
	suite.manager.UpdateResourceGroups(ctx, map[string]*rgpb.ResourceGroupConfig{
		DefaultResourceGroupName: newResourceGroupConfig(0, 0),
		"rg3":                    newResourceGroupConfig(10, 10),
		"rg2":                    newResourceGroupConfig(80, 80),
		"rg1":                    newResourceGroupConfig(10, 10),
	})
	suite.manager.AutoRecoverResourceGroup(ctx, DefaultResourceGroupName)
	suite.manager.AutoRecoverResourceGroup(ctx, "rg1")
	suite.manager.AutoRecoverResourceGroup(ctx, "rg2")
	suite.manager.AutoRecoverResourceGroup(ctx, "rg3")
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Equal(80, suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg3").NodeNum())
	suite.Equal(0, suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())

	suite.manager.UpdateResourceGroups(ctx, map[string]*rgpb.ResourceGroupConfig{
		DefaultResourceGroupName: newResourceGroupConfig(0, 5),
		"rg3":                    newResourceGroupConfig(5, 5),
		"rg2":                    newResourceGroupConfig(80, 80),
		"rg1":                    newResourceGroupConfig(20, 30),
	})
	suite.manager.AutoRecoverResourceGroup(ctx, "rg3") // recover redundant to missing rg.
	suite.Equal(15, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Equal(80, suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	suite.Equal(5, suite.manager.GetResourceGroup(ctx, "rg3").NodeNum())
	suite.Equal(0, suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())
	suite.manager.updateResourceGroups(ctx, map[string]*rgpb.ResourceGroupConfig{
		DefaultResourceGroupName: newResourceGroupConfig(5, 5),
		"rg3":                    newResourceGroupConfig(5, 10),
		"rg2":                    newResourceGroupConfig(80, 80),
		"rg1":                    newResourceGroupConfig(10, 10),
	})
	suite.manager.AutoRecoverResourceGroup(ctx, DefaultResourceGroupName) // recover missing from redundant rg.
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Equal(80, suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	suite.Equal(5, suite.manager.GetResourceGroup(ctx, "rg3").NodeNum())
	suite.Equal(5, suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())
}

func (suite *ResourceManagerSuite) testTransferNode() {
	ctx := suite.ctx
	// Test redundant nodes recover to default resource group.
	suite.manager.UpdateResourceGroups(ctx, map[string]*rgpb.ResourceGroupConfig{
		DefaultResourceGroupName: newResourceGroupConfig(40, 40),
		"rg3":                    newResourceGroupConfig(0, 0),
		"rg2":                    newResourceGroupConfig(40, 40),
		"rg1":                    newResourceGroupConfig(20, 20),
	})
	suite.manager.AutoRecoverResourceGroup(ctx, "rg1")
	suite.manager.AutoRecoverResourceGroup(ctx, "rg2")
	suite.manager.AutoRecoverResourceGroup(ctx, DefaultResourceGroupName)
	suite.manager.AutoRecoverResourceGroup(ctx, "rg3")

	suite.Equal(20, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Equal(40, suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	suite.Equal(0, suite.manager.GetResourceGroup(ctx, "rg3").NodeNum())
	suite.Equal(40, suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())

	// Test TransferNode.
	// param error.
	err := suite.manager.TransferNode(ctx, "rg1", "rg1", 1)
	suite.Error(err)

	err = suite.manager.TransferNode(ctx, "rg1", "rg2", 0)
	suite.Error(err)

	err = suite.manager.TransferNode(ctx, "rg3", "rg2", 1)
	suite.Error(err)

	err = suite.manager.TransferNode(ctx, "rg1", "rg10086", 1)
	suite.Error(err)

	err = suite.manager.TransferNode(ctx, "rg10086", "rg2", 1)
	suite.Error(err)

	// success
	err = suite.manager.TransferNode(ctx, "rg1", "rg3", 5)
	suite.NoError(err)

	suite.manager.AutoRecoverResourceGroup(ctx, "rg1")
	suite.manager.AutoRecoverResourceGroup(ctx, "rg2")
	suite.manager.AutoRecoverResourceGroup(ctx, DefaultResourceGroupName)
	suite.manager.AutoRecoverResourceGroup(ctx, "rg3")

	suite.Equal(15, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Equal(40, suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	suite.Equal(5, suite.manager.GetResourceGroup(ctx, "rg3").NodeNum())
	suite.Equal(40, suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())
}

func (suite *ResourceManagerSuite) TestIncomingNode() {
	ctx := suite.ctx
	suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   1,
		Address:  "localhost",
		Hostname: "localhost",
	}))
	suite.manager.incomingNode.Insert(1)

	suite.Equal(1, suite.manager.CheckIncomingNodeNum(ctx))
	suite.manager.AssignPendingIncomingNode(ctx)
	suite.Equal(0, suite.manager.CheckIncomingNodeNum(ctx))
	nodes, err := suite.manager.GetNodes(ctx, DefaultResourceGroupName)
	suite.NoError(err)
	suite.Len(nodes, 1)
}

func (suite *ResourceManagerSuite) TestUnassignFail() {
	ctx := suite.ctx
	// suite.man
	mockKV := mocks.NewMetaKv(suite.T())
	mockKV.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(nil).Once()

	store := querycoord.NewCatalog(mockKV)
	suite.manager = NewResourceManager(store, session.NewNodeManager())

	suite.manager.UpdateResourceGroups(ctx, map[string]*rgpb.ResourceGroupConfig{
		"rg1": newResourceGroupConfig(20, 30),
	})

	suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   1,
		Address:  "localhost",
		Hostname: "localhost",
	}))
	suite.manager.HandleNodeUp(ctx, 1)

	mockKV.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(merr.WrapErrServiceInternal("mocked")).Once()

	suite.Panics(func() {
		suite.manager.HandleNodeDown(ctx, 1)
	})
}

func TestGetResourceGroupsJSON(t *testing.T) {
	ctx := context.Background()
	nodeManager := session.NewNodeManager()
	manager := &ResourceManager{groups: make(map[string]*ResourceGroup)}
	rg1 := NewResourceGroup("rg1", newResourceGroupConfig(0, 10), nodeManager)
	rg1.nodes = typeutil.NewUniqueSet(1, 2)
	rg2 := NewResourceGroup("rg2", newResourceGroupConfig(0, 20), nodeManager)
	rg2.nodes = typeutil.NewUniqueSet(3, 4)
	manager.groups["rg1"] = rg1
	manager.groups["rg2"] = rg2

	jsonOutput := manager.GetResourceGroupsJSON(ctx)
	var resourceGroups []*metricsinfo.ResourceGroup
	err := json.Unmarshal([]byte(jsonOutput), &resourceGroups)
	assert.NoError(t, err)
	assert.Len(t, resourceGroups, 2)

	checkResult := func(rg *metricsinfo.ResourceGroup) {
		if rg.Name == "rg1" {
			assert.ElementsMatch(t, []int64{1, 2}, rg.Nodes)
		} else if rg.Name == "rg2" {
			assert.ElementsMatch(t, []int64{3, 4}, rg.Nodes)
		} else {
			assert.Failf(t, "unexpected resource group name", "unexpected resource group name %s", rg.Name)
		}
	}

	for _, rg := range resourceGroups {
		checkResult(rg)
	}
}

func (suite *ResourceManagerSuite) TestNodeLabels_NodeAssign() {
	ctx := suite.ctx
	suite.manager.AddResourceGroup(ctx, "rg1", &rgpb.ResourceGroupConfig{
		Requests: &rgpb.ResourceGroupLimit{
			NodeNum: 10,
		},
		Limits: &rgpb.ResourceGroupLimit{
			NodeNum: 10,
		},
		NodeFilter: &rgpb.ResourceGroupNodeFilter{
			NodeLabels: []*commonpb.KeyValuePair{
				{
					Key:   "dc_name",
					Value: "label1",
				},
			},
		},
	})

	suite.manager.AddResourceGroup(ctx, "rg2", &rgpb.ResourceGroupConfig{
		Requests: &rgpb.ResourceGroupLimit{
			NodeNum: 10,
		},
		Limits: &rgpb.ResourceGroupLimit{
			NodeNum: 10,
		},
		NodeFilter: &rgpb.ResourceGroupNodeFilter{
			NodeLabels: []*commonpb.KeyValuePair{
				{
					Key:   "dc_name",
					Value: "label2",
				},
			},
		},
	})

	suite.manager.AddResourceGroup(ctx, "rg3", &rgpb.ResourceGroupConfig{
		Requests: &rgpb.ResourceGroupLimit{
			NodeNum: 10,
		},
		Limits: &rgpb.ResourceGroupLimit{
			NodeNum: 10,
		},
		NodeFilter: &rgpb.ResourceGroupNodeFilter{
			NodeLabels: []*commonpb.KeyValuePair{
				{
					Key:   "dc_name",
					Value: "label3",
				},
			},
		},
	})

	// test that all query nodes has been marked label1
	for i := 1; i <= 30; i++ {
		suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
			NodeID:   int64(i),
			Address:  "localhost",
			Hostname: "localhost",
			Labels: map[string]string{
				"dc_name": "label1",
			},
		}))
		suite.manager.HandleNodeUp(ctx, int64(i))
	}
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Equal(0, suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	suite.Equal(0, suite.manager.GetResourceGroup(ctx, "rg3").NodeNum())
	suite.Equal(20, suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())

	// test new querynode with label2
	for i := 31; i <= 40; i++ {
		suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
			NodeID:   int64(i),
			Address:  "localhost",
			Hostname: "localhost",
			Labels: map[string]string{
				"dc_name": "label2",
			},
		}))
		suite.manager.HandleNodeUp(ctx, int64(i))
	}
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	suite.Equal(0, suite.manager.GetResourceGroup(ctx, "rg3").NodeNum())
	suite.Equal(20, suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())
	nodesInRG, _ := suite.manager.GetNodes(ctx, "rg2")
	for _, node := range nodesInRG {
		suite.Equal("label2", suite.manager.nodeMgr.Get(node).Labels()["dc_name"])
	}

	// test new querynode with label3
	for i := 41; i <= 50; i++ {
		suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
			NodeID:   int64(i),
			Address:  "localhost",
			Hostname: "localhost",
			Labels: map[string]string{
				"dc_name": "label3",
			},
		}))
		suite.manager.HandleNodeUp(ctx, int64(i))
	}
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg3").NodeNum())
	suite.Equal(20, suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())
	nodesInRG, _ = suite.manager.GetNodes(ctx, "rg3")
	for _, node := range nodesInRG {
		suite.Equal("label3", suite.manager.nodeMgr.Get(node).Labels()["dc_name"])
	}

	// test swap rg's label
	suite.manager.UpdateResourceGroups(ctx, map[string]*rgpb.ResourceGroupConfig{
		"rg1": {
			Requests: &rgpb.ResourceGroupLimit{
				NodeNum: 10,
			},
			Limits: &rgpb.ResourceGroupLimit{
				NodeNum: 10,
			},
			NodeFilter: &rgpb.ResourceGroupNodeFilter{
				NodeLabels: []*commonpb.KeyValuePair{
					{
						Key:   "dc_name",
						Value: "label2",
					},
				},
			},
		},
		"rg2": {
			Requests: &rgpb.ResourceGroupLimit{
				NodeNum: 10,
			},
			Limits: &rgpb.ResourceGroupLimit{
				NodeNum: 10,
			},
			NodeFilter: &rgpb.ResourceGroupNodeFilter{
				NodeLabels: []*commonpb.KeyValuePair{
					{
						Key:   "dc_name",
						Value: "label3",
					},
				},
			},
		},
		"rg3": {
			Requests: &rgpb.ResourceGroupLimit{
				NodeNum: 10,
			},
			Limits: &rgpb.ResourceGroupLimit{
				NodeNum: 10,
			},
			NodeFilter: &rgpb.ResourceGroupNodeFilter{
				NodeLabels: []*commonpb.KeyValuePair{
					{
						Key:   "dc_name",
						Value: "label1",
					},
				},
			},
		},
	})

	log.Info("test swap rg's label")
	for i := 0; i < 4; i++ {
		suite.manager.AutoRecoverResourceGroup(ctx, "rg1")
		suite.manager.AutoRecoverResourceGroup(ctx, "rg2")
		suite.manager.AutoRecoverResourceGroup(ctx, "rg3")
		suite.manager.AutoRecoverResourceGroup(ctx, DefaultResourceGroupName)
	}
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg3").NodeNum())
	suite.Equal(20, suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())
	nodesInRG, _ = suite.manager.GetNodes(ctx, "rg1")
	for _, node := range nodesInRG {
		suite.Equal("label2", suite.manager.nodeMgr.Get(node).Labels()["dc_name"])
	}

	nodesInRG, _ = suite.manager.GetNodes(ctx, "rg2")
	for _, node := range nodesInRG {
		suite.Equal("label3", suite.manager.nodeMgr.Get(node).Labels()["dc_name"])
	}

	nodesInRG, _ = suite.manager.GetNodes(ctx, "rg3")
	for _, node := range nodesInRG {
		suite.Equal("label1", suite.manager.nodeMgr.Get(node).Labels()["dc_name"])
	}
}

func (suite *ResourceManagerSuite) TestNodeLabels_NodeDown() {
	ctx := suite.ctx
	suite.manager.AddResourceGroup(ctx, "rg1", &rgpb.ResourceGroupConfig{
		Requests: &rgpb.ResourceGroupLimit{
			NodeNum: 10,
		},
		Limits: &rgpb.ResourceGroupLimit{
			NodeNum: 10,
		},
		NodeFilter: &rgpb.ResourceGroupNodeFilter{
			NodeLabels: []*commonpb.KeyValuePair{
				{
					Key:   "dc_name",
					Value: "label1",
				},
			},
		},
	})

	suite.manager.AddResourceGroup(ctx, "rg2", &rgpb.ResourceGroupConfig{
		Requests: &rgpb.ResourceGroupLimit{
			NodeNum: 10,
		},
		Limits: &rgpb.ResourceGroupLimit{
			NodeNum: 10,
		},
		NodeFilter: &rgpb.ResourceGroupNodeFilter{
			NodeLabels: []*commonpb.KeyValuePair{
				{
					Key:   "dc_name",
					Value: "label2",
				},
			},
		},
	})

	suite.manager.AddResourceGroup(ctx, "rg3", &rgpb.ResourceGroupConfig{
		Requests: &rgpb.ResourceGroupLimit{
			NodeNum: 10,
		},
		Limits: &rgpb.ResourceGroupLimit{
			NodeNum: 10,
		},
		NodeFilter: &rgpb.ResourceGroupNodeFilter{
			NodeLabels: []*commonpb.KeyValuePair{
				{
					Key:   "dc_name",
					Value: "label3",
				},
			},
		},
	})

	// test that all query nodes has been marked label1
	for i := 1; i <= 10; i++ {
		suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
			NodeID:   int64(i),
			Address:  "localhost",
			Hostname: "localhost",
			Labels: map[string]string{
				"dc_name": "label1",
			},
		}))
		suite.manager.HandleNodeUp(ctx, int64(i))
	}

	// test new querynode with label2
	for i := 31; i <= 40; i++ {
		suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
			NodeID:   int64(i),
			Address:  "localhost",
			Hostname: "localhost",
			Labels: map[string]string{
				"dc_name": "label2",
			},
		}))
		suite.manager.HandleNodeUp(ctx, int64(i))
	}
	// test new querynode with label3
	for i := 41; i <= 50; i++ {
		suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
			NodeID:   int64(i),
			Address:  "localhost",
			Hostname: "localhost",
			Labels: map[string]string{
				"dc_name": "label3",
			},
		}))
		suite.manager.HandleNodeUp(ctx, int64(i))
	}
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg3").NodeNum())

	// test node down with label1
	suite.manager.HandleNodeDown(ctx, int64(1))
	suite.manager.nodeMgr.Remove(int64(1))
	suite.Equal(9, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg3").NodeNum())

	// test node up with label2
	suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   int64(101),
		Address:  "localhost",
		Hostname: "localhost",
		Labels: map[string]string{
			"dc_name": "label2",
		},
	}))
	suite.manager.HandleNodeUp(ctx, int64(101))
	suite.Equal(9, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg3").NodeNum())
	suite.Equal(1, suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())

	// test node up with label1
	suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   int64(102),
		Address:  "localhost",
		Hostname: "localhost",
		Labels: map[string]string{
			"dc_name": "label1",
		},
	}))
	suite.manager.HandleNodeUp(ctx, int64(102))
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg1").NodeNum())
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg2").NodeNum())
	suite.Equal(10, suite.manager.GetResourceGroup(ctx, "rg3").NodeNum())
	suite.Equal(1, suite.manager.GetResourceGroup(ctx, DefaultResourceGroupName).NodeNum())
	nodesInRG, _ := suite.manager.GetNodes(ctx, "rg1")
	for _, node := range nodesInRG {
		suite.Equal("label1", suite.manager.nodeMgr.Get(node).Labels()["dc_name"])
	}

	suite.manager.AutoRecoverResourceGroup(ctx, "rg1")
	suite.manager.AutoRecoverResourceGroup(ctx, "rg2")
	suite.manager.AutoRecoverResourceGroup(ctx, "rg3")
	suite.manager.AutoRecoverResourceGroup(ctx, DefaultResourceGroupName)
	nodesInRG, _ = suite.manager.GetNodes(ctx, DefaultResourceGroupName)
	for _, node := range nodesInRG {
		suite.Equal("label2", suite.manager.nodeMgr.Get(node).Labels()["dc_name"])
	}
}

// createTestResourceManager creates a ResourceManager for testing
func createTestResourceManager(t *testing.T) *ResourceManager {
	// Create a mock catalog
	mockCatalog := &mocks.MetaKv{}
	mockCatalog.On("MultiSave", mock.Anything, mock.Anything).Return(nil)

	// Create a mock node manager
	nodeMgr := session.NewNodeManager()

	// Create resource manager
	store := querycoord.NewCatalog(mockCatalog)
	manager := NewResourceManager(store, nodeMgr)

	return manager
}

// TestResourceManager_handleNodeUp tests the private handleNodeUp method
func TestResourceManager_handleNodeUp(t *testing.T) {
	// Arrange
	manager := createTestResourceManager(t)
	ctx := context.Background()
	nodeID := int64(1001)

	// Add node to node manager
	nodeInfo := session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   nodeID,
		Address:  "localhost",
		Hostname: "localhost",
	})
	manager.nodeMgr.Add(nodeInfo)

	// Act
	manager.handleNodeUp(ctx, nodeID)

	// Assert
	// After successful assignment, node should be removed from incomingNode
	assert.False(t, manager.incomingNode.Contain(nodeID))

	// Verify node was assigned to default resource group
	nodes, err := manager.GetNodes(ctx, DefaultResourceGroupName)
	assert.NoError(t, err)
	assert.Contains(t, nodes, nodeID)
}

// TestResourceManager_handleNodeDown tests the private handleNodeDown method
func TestResourceManager_handleNodeDown(t *testing.T) {
	// Arrange
	manager := createTestResourceManager(t)
	ctx := context.Background()
	nodeID := int64(1002)

	// Add node to node manager
	nodeInfo := session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   nodeID,
		Address:  "localhost",
		Hostname: "localhost",
	})
	manager.nodeMgr.Add(nodeInfo)

	// Add node to incoming set and assign it to a resource group first
	manager.handleNodeUp(ctx, nodeID)
	nodes, err := manager.GetNodes(ctx, DefaultResourceGroupName)
	assert.NoError(t, err)
	assert.Contains(t, nodes, nodeID)

	// Act
	manager.handleNodeDown(ctx, nodeID)

	// Assert
	assert.False(t, manager.incomingNode.Contain(nodeID))

	// Verify node was removed from resource group
	nodes, err = manager.GetNodes(ctx, DefaultResourceGroupName)
	assert.NoError(t, err)
	assert.NotContains(t, nodes, nodeID)

	// Verify node is no longer in nodeIDMap
	_, exists := manager.nodeIDMap[nodeID]
	assert.False(t, exists)
}

// TestResourceManager_handleNodeStopping tests the private handleNodeStopping method
func TestResourceManager_handleNodeStopping(t *testing.T) {
	// Arrange
	manager := createTestResourceManager(t)
	ctx := context.Background()
	nodeID := int64(1003)

	// Add node to node manager
	nodeInfo := session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   nodeID,
		Address:  "localhost",
		Hostname: "localhost",
	})
	manager.nodeMgr.Add(nodeInfo)

	// Add node to incoming set and assign it to a resource group first
	manager.handleNodeUp(ctx, nodeID)
	nodes, err := manager.GetNodes(ctx, DefaultResourceGroupName)
	assert.NoError(t, err)
	assert.Contains(t, nodes, nodeID)

	// Act
	manager.handleNodeStopping(ctx, nodeID)

	// Assert
	assert.False(t, manager.incomingNode.Contain(nodeID))

	// Verify node was removed from resource group
	nodes, err = manager.GetNodes(ctx, DefaultResourceGroupName)
	assert.NoError(t, err)
	assert.NotContains(t, nodes, nodeID)

	// Verify node is no longer in nodeIDMap
	_, exists := manager.nodeIDMap[nodeID]
	assert.False(t, exists)
}

// TestResourceManager_CheckNodesInResourceGroup tests the CheckNodesInResourceGroup method
func TestResourceManager_CheckNodesInResourceGroup(t *testing.T) {
	// Arrange
	manager := createTestResourceManager(t)

	// Add some nodes to node manager
	nodeInfo1 := session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   1001,
		Address:  "localhost:1001",
		Hostname: "localhost",
	})
	nodeInfo2 := session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   1002,
		Address:  "localhost:1002",
		Hostname: "localhost",
	})
	nodeInfo3 := session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   1003,
		Address:  "localhost:1003",
		Hostname: "localhost",
	})
	manager.nodeMgr.Add(nodeInfo1)
	manager.nodeMgr.Add(nodeInfo2)
	manager.nodeMgr.Add(nodeInfo3)

	// Set node 1002 as stopping
	nodeInfo2.SetState(session.NodeStateStopping)

	// Add nodes to default resource group
	ctx := context.Background()
	manager.handleNodeUp(ctx, 1001)
	manager.handleNodeUp(ctx, 1002)
	manager.handleNodeUp(ctx, 1004)

	// Act
	manager.CheckNodesInResourceGroup(ctx)

	// Verify final state: offline node (1004) should be removed
	finalNodes, err := manager.GetNodes(context.Background(), DefaultResourceGroupName)
	assert.NoError(t, err)
	assert.NotContains(t, finalNodes, int64(1004), "Offline node should be removed")

	// Verify stopping node (1002) should be removed
	assert.NotContains(t, finalNodes, int64(1002), "Stopping node should be removed")

	// Verify healthy node (1001) should remain
	assert.Contains(t, finalNodes, int64(1001), "Healthy node should remain")

	// Verify new node (1003) should be added
	assert.Contains(t, finalNodes, int64(1003), "New node should be added")
}

// TestResourceManager_CheckNodesInResourceGroup_AllNodesHealthy tests CheckNodesInResourceGroup with all healthy nodes
func TestResourceManager_CheckNodesInResourceGroup_AllNodesHealthy(t *testing.T) {
	// Arrange
	manager := createTestResourceManager(t)

	// Add some healthy nodes to node manager
	nodeInfo1 := session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   1001,
		Address:  "localhost:1001",
		Hostname: "localhost",
	})
	nodeInfo2 := session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   1002,
		Address:  "localhost:1002",
		Hostname: "localhost",
	})
	manager.nodeMgr.Add(nodeInfo1)
	manager.nodeMgr.Add(nodeInfo2)

	// Add nodes to default resource group
	ctx := context.Background()
	manager.handleNodeUp(ctx, 1001)
	manager.handleNodeUp(ctx, 1002)

	// Act
	manager.CheckNodesInResourceGroup(ctx)

	// Verify that healthy nodes remain unchanged
	finalNodes, err := manager.GetNodes(ctx, DefaultResourceGroupName)
	assert.NoError(t, err)
	assert.Contains(t, finalNodes, int64(1001), "Healthy node should remain")
	assert.Contains(t, finalNodes, int64(1002), "Healthy node should remain")
	assert.Equal(t, 2, len(finalNodes), "Should have exactly 2 nodes")
}