github.com/cloudreve/Cloudreve/v3@v3.0.0-20240224133659-3edb00a6484c/pkg/cluster/pool_test.go (about) 1 package cluster 2 3 import ( 4 "database/sql" 5 "errors" 6 "github.com/DATA-DOG/go-sqlmock" 7 model "github.com/cloudreve/Cloudreve/v3/models" 8 "github.com/cloudreve/Cloudreve/v3/pkg/balancer" 9 "github.com/jinzhu/gorm" 10 "github.com/stretchr/testify/assert" 11 "testing" 12 ) 13 14 var mock sqlmock.Sqlmock 15 16 // TestMain 初始化数据库Mock 17 func TestMain(m *testing.M) { 18 var db *sql.DB 19 var err error 20 db, mock, err = sqlmock.New() 21 if err != nil { 22 panic("An error was not expected when opening a stub database connection") 23 } 24 model.DB, _ = gorm.Open("mysql", db) 25 defer db.Close() 26 m.Run() 27 } 28 29 func TestInitFailed(t *testing.T) { 30 a := assert.New(t) 31 mock.ExpectQuery("SELECT(.+)").WillReturnError(errors.New("error")) 32 Init() 33 a.NoError(mock.ExpectationsWereMet()) 34 } 35 36 func TestInitSuccess(t *testing.T) { 37 a := assert.New(t) 38 mock.ExpectQuery("SELECT(.+)").WillReturnRows(sqlmock.NewRows([]string{"id", "aria2_enabled", "type"}).AddRow(1, true, model.MasterNodeType)) 39 Init() 40 a.NoError(mock.ExpectationsWereMet()) 41 } 42 43 func TestNodePool_GetNodeByID(t *testing.T) { 44 a := assert.New(t) 45 p := &NodePool{} 46 p.Init() 47 mockNode := &nodeMock{} 48 49 // inactive 50 { 51 p.inactive[1] = mockNode 52 a.Equal(mockNode, p.GetNodeByID(1)) 53 } 54 55 // active 56 { 57 delete(p.inactive, 1) 58 p.active[1] = mockNode 59 a.Equal(mockNode, p.GetNodeByID(1)) 60 } 61 } 62 63 func TestNodePool_NodeStatusChange(t *testing.T) { 64 a := assert.New(t) 65 p := &NodePool{} 66 n := &MasterNode{Model: &model.Node{}} 67 p.Init() 68 p.inactive[1] = n 69 70 p.nodeStatusChange(true, 1) 71 a.Len(p.inactive, 0) 72 a.Equal(n, p.active[1]) 73 74 p.nodeStatusChange(false, 1) 75 a.Len(p.active, 0) 76 a.Equal(n, p.inactive[1]) 77 78 p.nodeStatusChange(false, 1) 79 a.Len(p.active, 0) 80 a.Equal(n, p.inactive[1]) 81 } 82 83 func TestNodePool_Add(t *testing.T) { 84 a := assert.New(t) 85 p := &NodePool{} 86 p.Init() 87 88 // new node 89 { 90 p.Add(&model.Node{}) 91 a.Len(p.active, 1) 92 } 93 94 // old node 95 { 96 p.inactive[0] = p.active[0] 97 delete(p.active, 0) 98 p.Add(&model.Node{}) 99 a.Len(p.active, 0) 100 a.Len(p.inactive, 1) 101 } 102 } 103 104 func TestNodePool_Delete(t *testing.T) { 105 a := assert.New(t) 106 p := &NodePool{} 107 p.Init() 108 109 // active 110 { 111 mockNode := &nodeMock{} 112 mockNode.On("Kill") 113 p.active[0] = mockNode 114 p.Delete(0) 115 a.Len(p.active, 0) 116 a.Len(p.inactive, 0) 117 mockNode.AssertExpectations(t) 118 } 119 120 p.Init() 121 122 // inactive 123 { 124 mockNode := &nodeMock{} 125 mockNode.On("Kill") 126 p.inactive[0] = mockNode 127 p.Delete(0) 128 a.Len(p.active, 0) 129 a.Len(p.inactive, 0) 130 mockNode.AssertExpectations(t) 131 } 132 } 133 134 func TestNodePool_BalanceNodeByFeature(t *testing.T) { 135 a := assert.New(t) 136 p := &NodePool{} 137 p.Init() 138 139 // success 140 { 141 p.featureMap["test"] = []Node{&MasterNode{}} 142 err, res := p.BalanceNodeByFeature("test", balancer.NewBalancer("round-robin")) 143 a.NoError(err) 144 a.Equal(p.featureMap["test"][0], res) 145 } 146 147 // NoNodes 148 { 149 p.featureMap["test"] = []Node{} 150 err, res := p.BalanceNodeByFeature("test", balancer.NewBalancer("round-robin")) 151 a.Error(err) 152 a.Nil(res) 153 } 154 155 // No match feature 156 { 157 err, res := p.BalanceNodeByFeature("test2", balancer.NewBalancer("round-robin")) 158 a.Error(err) 159 a.Nil(res) 160 } 161 }