Test: pkg/cluster

pkg/cluster/pool_test.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"github.com/DATA-DOG/go-sqlmock"
 	model "github.com/cloudreve/Cloudreve/v3/models"
+	"github.com/cloudreve/Cloudreve/v3/pkg/balancer"
 	"github.com/jinzhu/gorm"
 	"github.com/stretchr/testify/assert"
 	"testing"
@@ -60,5 +61,101 @@ func TestNodePool_GetNodeByID(t *testing.T) {
 }
 
 func TestNodePool_NodeStatusChange(t *testing.T) {
+	a := assert.New(t)
+	p := &NodePool{}
+	n := &MasterNode{Model: &model.Node{}}
+	p.Init()
+	p.inactive[1] = n
 
+	p.nodeStatusChange(true, 1)
+	a.Len(p.inactive, 0)
+	a.Equal(n, p.active[1])
+
+	p.nodeStatusChange(false, 1)
+	a.Len(p.active, 0)
+	a.Equal(n, p.inactive[1])
+
+	p.nodeStatusChange(false, 1)
+	a.Len(p.active, 0)
+	a.Equal(n, p.inactive[1])
+}
+
+func TestNodePool_Add(t *testing.T) {
+	a := assert.New(t)
+	p := &NodePool{}
+	p.Init()
+
+	// new node
+	{
+		p.Add(&model.Node{})
+		a.Len(p.active, 1)
+	}
+
+	// old node
+	{
+		p.inactive[0] = p.active[0]
+		delete(p.active, 0)
+		p.Add(&model.Node{})
+		a.Len(p.active, 0)
+		a.Len(p.inactive, 1)
+	}
+}
+
+func TestNodePool_Delete(t *testing.T) {
+	a := assert.New(t)
+	p := &NodePool{}
+	p.Init()
+
+	// active
+	{
+		mockNode := &nodeMock{}
+		mockNode.On("Kill")
+		p.active[0] = mockNode
+		p.Delete(0)
+		a.Len(p.active, 0)
+		a.Len(p.inactive, 0)
+		mockNode.AssertExpectations(t)
+	}
+
+	p.Init()
+
+	// inactive
+	{
+		mockNode := &nodeMock{}
+		mockNode.On("Kill")
+		p.inactive[0] = mockNode
+		p.Delete(0)
+		a.Len(p.active, 0)
+		a.Len(p.inactive, 0)
+		mockNode.AssertExpectations(t)
+	}
+}
+
+func TestNodePool_BalanceNodeByFeature(t *testing.T) {
+	a := assert.New(t)
+	p := &NodePool{}
+	p.Init()
+
+	// success
+	{
+		p.featureMap["test"] = []Node{&MasterNode{}}
+		err, res := p.BalanceNodeByFeature("test", balancer.NewBalancer("round-robin"))
+		a.NoError(err)
+		a.Equal(p.featureMap["test"][0], res)
+	}
+
+	// NoNodes
+	{
+		p.featureMap["test"] = []Node{}
+		err, res := p.BalanceNodeByFeature("test", balancer.NewBalancer("round-robin"))
+		a.Error(err)
+		a.Nil(res)
+	}
+
+	// No match feature
+	{
+		err, res := p.BalanceNodeByFeature("test2", balancer.NewBalancer("round-robin"))
+		a.Error(err)
+		a.Nil(res)
+	}
 }