kikoqiu
4 years ago
committed by
GitHub
commit
4810756add
@ -0,0 +1,156 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"errors"
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/aria2/common"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/auth"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/mq"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
"github.com/stretchr/testify/assert"
|
||||
testMock "github.com/stretchr/testify/mock"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestMasterMetadata(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
masterMetaDataFunc := MasterMetadata()
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest("GET", "/", nil)
|
||||
|
||||
c.Request.Header = map[string][]string{
|
||||
"X-Site-Id": {"expectedSiteID"},
|
||||
"X-Site-Url": {"expectedSiteURL"},
|
||||
"X-Cloudreve-Version": {"expectedMasterVersion"},
|
||||
}
|
||||
masterMetaDataFunc(c)
|
||||
siteID, _ := c.Get("MasterSiteID")
|
||||
siteURL, _ := c.Get("MasterSiteURL")
|
||||
siteVersion, _ := c.Get("MasterVersion")
|
||||
|
||||
a.Equal("expectedSiteID", siteID.(string))
|
||||
a.Equal("expectedSiteURL", siteURL.(string))
|
||||
a.Equal("expectedMasterVersion", siteVersion.(string))
|
||||
}
|
||||
|
||||
func TestSlaveRPCSignRequired(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
np := &cluster.NodePool{}
|
||||
np.Init()
|
||||
slaveRPCSignRequiredFunc := SlaveRPCSignRequired(np)
|
||||
rec := httptest.NewRecorder()
|
||||
|
||||
// id parse failed
|
||||
{
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest("GET", "/", nil)
|
||||
c.Request.Header.Set("X-Node-Id", "unknown")
|
||||
slaveRPCSignRequiredFunc(c)
|
||||
a.True(c.IsAborted())
|
||||
}
|
||||
|
||||
// node id not exist
|
||||
{
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest("GET", "/", nil)
|
||||
c.Request.Header.Set("X-Node-Id", "38")
|
||||
slaveRPCSignRequiredFunc(c)
|
||||
a.True(c.IsAborted())
|
||||
}
|
||||
|
||||
// success
|
||||
{
|
||||
authInstance := auth.HMACAuth{SecretKey: []byte("")}
|
||||
np.Add(&model.Node{Model: gorm.Model{
|
||||
ID: 38,
|
||||
}})
|
||||
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest("POST", "/", nil)
|
||||
c.Request.Header.Set("X-Node-Id", "38")
|
||||
c.Request = auth.SignRequest(authInstance, c.Request, 0)
|
||||
slaveRPCSignRequiredFunc(c)
|
||||
a.False(c.IsAborted())
|
||||
}
|
||||
}
|
||||
|
||||
type SlaveControllerMock struct {
|
||||
testMock.Mock
|
||||
}
|
||||
|
||||
func (s SlaveControllerMock) HandleHeartBeat(pingReq *serializer.NodePingReq) (serializer.NodePingResp, error) {
|
||||
args := s.Called(pingReq)
|
||||
return args.Get(0).(serializer.NodePingResp), args.Error(1)
|
||||
}
|
||||
|
||||
func (s SlaveControllerMock) GetAria2Instance(s2 string) (common.Aria2, error) {
|
||||
args := s.Called(s2)
|
||||
return args.Get(0).(common.Aria2), args.Error(1)
|
||||
}
|
||||
|
||||
func (s SlaveControllerMock) SendNotification(s3 string, s2 string, message mq.Message) error {
|
||||
args := s.Called(s3, s2, message)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (s SlaveControllerMock) SubmitTask(s3 string, i interface{}, s2 string, f func(interface{})) error {
|
||||
args := s.Called(s3, i, s2, f)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (s SlaveControllerMock) GetMasterInfo(s2 string) (*cluster.MasterInfo, error) {
|
||||
args := s.Called(s2)
|
||||
return args.Get(0).(*cluster.MasterInfo), args.Error(1)
|
||||
}
|
||||
|
||||
func (s SlaveControllerMock) GetOneDriveToken(s2 string, u uint) (string, error) {
|
||||
args := s.Called(s2, u)
|
||||
return args.String(0), args.Error(1)
|
||||
}
|
||||
|
||||
func TestUseSlaveAria2Instance(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
|
||||
// MasterSiteID not set
|
||||
{
|
||||
testController := &SlaveControllerMock{}
|
||||
useSlaveAria2InstanceFunc := UseSlaveAria2Instance(testController)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest("GET", "/", nil)
|
||||
useSlaveAria2InstanceFunc(c)
|
||||
a.True(c.IsAborted())
|
||||
}
|
||||
|
||||
// Cannot get aria2 instances
|
||||
{
|
||||
testController := &SlaveControllerMock{}
|
||||
useSlaveAria2InstanceFunc := UseSlaveAria2Instance(testController)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest("GET", "/", nil)
|
||||
c.Set("MasterSiteID", "expectedSiteID")
|
||||
testController.On("GetAria2Instance", "expectedSiteID").Return(&common.DummyAria2{}, errors.New("error"))
|
||||
useSlaveAria2InstanceFunc(c)
|
||||
a.True(c.IsAborted())
|
||||
testController.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// Success
|
||||
{
|
||||
testController := &SlaveControllerMock{}
|
||||
useSlaveAria2InstanceFunc := UseSlaveAria2Instance(testController)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest("GET", "/", nil)
|
||||
c.Set("MasterSiteID", "expectedSiteID")
|
||||
testController.On("GetAria2Instance", "expectedSiteID").Return(&common.DummyAria2{}, nil)
|
||||
useSlaveAria2InstanceFunc(c)
|
||||
a.False(c.IsAborted())
|
||||
res, _ := c.Get("MasterAria2Instance")
|
||||
a.NotNil(res)
|
||||
testController.AssertExpectations(t)
|
||||
}
|
||||
}
|
||||
@ -0,0 +1,64 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestGetNodeByID(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
mock.ExpectQuery("SELECT(.+)nodes").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
res, err := GetNodeByID(1)
|
||||
a.NoError(err)
|
||||
a.EqualValues(1, res.ID)
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestGetNodesByStatus(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
mock.ExpectQuery("SELECT(.+)nodes").WillReturnRows(sqlmock.NewRows([]string{"status"}).AddRow(NodeActive))
|
||||
res, err := GetNodesByStatus(NodeActive)
|
||||
a.NoError(err)
|
||||
a.Len(res, 1)
|
||||
a.EqualValues(NodeActive, res[0].Status)
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestNode_AfterFind(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
node := &Node{}
|
||||
|
||||
// No aria2 options
|
||||
{
|
||||
a.NoError(node.AfterFind())
|
||||
}
|
||||
|
||||
// with aria2 options
|
||||
{
|
||||
node.Aria2Options = `{"timeout":1}`
|
||||
a.NoError(node.AfterFind())
|
||||
a.Equal(1, node.Aria2OptionsSerialized.Timeout)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNode_BeforeSave(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
node := &Node{}
|
||||
|
||||
node.Aria2OptionsSerialized.Timeout = 1
|
||||
a.NoError(node.BeforeSave())
|
||||
a.Contains("1", node.Aria2Options)
|
||||
}
|
||||
|
||||
func TestNode_SetStatus(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
node := &Node{}
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)nodes").WithArgs(NodeActive, sqlmock.AnyArg()).WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
a.NoError(node.SetStatus(NodeActive))
|
||||
a.Equal(NodeActive, node.Status)
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
}
|
||||
@ -1,8 +1,12 @@
|
||||
package cluster
|
||||
|
||||
import "errors"
|
||||
import (
|
||||
"errors"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrFeatureNotExist = errors.New("No nodes in nodepool match the feature specificed")
|
||||
ErrIlegalPath = errors.New("path out of boundary of setting temp folder")
|
||||
ErrMasterNotFound = serializer.NewError(serializer.CodeMasterNotFound, "未知的主机节点", nil)
|
||||
)
|
||||
|
||||
@ -1,7 +0,0 @@
|
||||
package slave
|
||||
|
||||
import "github.com/cloudreve/Cloudreve/v3/pkg/serializer"
|
||||
|
||||
var (
|
||||
ErrMasterNotFound = serializer.NewError(serializer.CodeMasterNotFound, "未知的主机节点", nil)
|
||||
)
|
||||
Loading…
Reference in new issue