Test: cluster/node.go and controller.go

pkg/cluster/pool_test.go

@@ -0,0 +1,64 @@
+package cluster
+
+import (
+	"database/sql"
+	"errors"
+	"github.com/DATA-DOG/go-sqlmock"
+	model "github.com/cloudreve/Cloudreve/v3/models"
+	"github.com/jinzhu/gorm"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+var mock sqlmock.Sqlmock
+
+// TestMain 初始化数据库Mock
+func TestMain(m *testing.M) {
+	var db *sql.DB
+	var err error
+	db, mock, err = sqlmock.New()
+	if err != nil {
+		panic("An error was not expected when opening a stub database connection")
+	}
+	model.DB, _ = gorm.Open("mysql", db)
+	defer db.Close()
+	m.Run()
+}
+
+func TestInitFailed(t *testing.T) {
+	a := assert.New(t)
+	mock.ExpectQuery("SELECT(.+)").WillReturnError(errors.New("error"))
+	Init()
+	a.NoError(mock.ExpectationsWereMet())
+}
+
+func TestInitSuccess(t *testing.T) {
+	a := assert.New(t)
+	mock.ExpectQuery("SELECT(.+)").WillReturnRows(sqlmock.NewRows([]string{"id", "aria2_enabled", "type"}).AddRow(1, true, model.MasterNodeType))
+	Init()
+	a.NoError(mock.ExpectationsWereMet())
+}
+
+func TestNodePool_GetNodeByID(t *testing.T) {
+	a := assert.New(t)
+	p := &NodePool{}
+	p.Init()
+	mockNode := &nodeMock{}
+
+	// inactive
+	{
+		p.inactive[1] = mockNode
+		a.Equal(mockNode, p.GetNodeByID(1))
+	}
+
+	// active
+	{
+		delete(p.inactive, 1)
+		p.active[1] = mockNode
+		a.Equal(mockNode, p.GetNodeByID(1))
+	}
+}
+
+func TestNodePool_NodeStatusChange(t *testing.T) {
+
+}