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  }