mirror of
https://github.com/halejohn/Cloudreve.git
synced 2026-01-26 09:34:57 +08:00
Feat: aria2 download and transfer in slave node (#1040)
* Feat: retrieve nodes from data table * Feat: master node ping slave node in REST API * Feat: master send scheduled ping request * Feat: inactive nodes recover loop * Modify: remove database operations from aria2 RPC caller implementation * Feat: init aria2 client in master node * Feat: Round Robin load balancer * Feat: create and monitor aria2 task in master node * Feat: salve receive and handle heartbeat * Fix: Node ID will be 0 in download record generated in older version * Feat: sign request headers with all `X-` prefix * Feat: API call to slave node will carry meta data in headers * Feat: call slave aria2 rpc method from master * Feat: get slave aria2 task status Feat: encode slave response data using gob * Feat: aria2 callback to master node / cancel or select task to slave node * Fix: use dummy aria2 client when caller initialize failed in master node * Feat: slave aria2 status event callback / salve RPC auth * Feat: prototype for slave driven filesystem * Feat: retry for init aria2 client in master node * Feat: init request client with global options * Feat: slave receive async task from master * Fix: competition write in request header * Refactor: dependency initialize order * Feat: generic message queue implementation * Feat: message queue implementation * Feat: master waiting slave transfer result * Feat: slave transfer file in stateless policy * Feat: slave transfer file in slave policy * Feat: slave transfer file in local policy * Feat: slave transfer file in OneDrive policy * Fix: failed to initialize update checker http client * Feat: list slave nodes for dashboard * Feat: test aria2 rpc connection in slave * Feat: add and save node * Feat: add and delete node in node pool * Fix: temp file cannot be removed when aria2 task fails * Fix: delete node in admin panel * Feat: edit node and get node info * Modify: delete unused settings
This commit is contained in:
AaronLiu
GitHub
279
pkg/aria2/monitor/monitor.go
Normal file
279
pkg/aria2/monitor/monitor.go
Normal file
@@ -0,0 +1,279 @@
|
||||
package monitor
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/aria2/common"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/aria2/rpc"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/filesystem"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/filesystem/driver/local"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/filesystem/fsctx"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/mq"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/task"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/util"
|
||||
)
|
||||
|
||||
// Monitor 离线下载状态监控
|
||||
type Monitor struct {
|
||||
Task *model.Download
|
||||
Interval time.Duration
|
||||
|
||||
notifier <-chan mq.Message
|
||||
node cluster.Node
|
||||
retried int
|
||||
}
|
||||
|
||||
var MAX_RETRY = 10
|
||||
|
||||
// NewMonitor 新建离线下载状态监控
|
||||
func NewMonitor(task *model.Download) {
|
||||
monitor := &Monitor{
|
||||
Task: task,
|
||||
notifier: make(chan mq.Message),
|
||||
node: cluster.Default.GetNodeByID(task.GetNodeID()),
|
||||
}
|
||||
|
||||
if monitor.node != nil {
|
||||
monitor.Interval = time.Duration(monitor.node.GetAria2Instance().GetConfig().Interval) * time.Second
|
||||
go monitor.Loop()
|
||||
|
||||
monitor.notifier = mq.GlobalMQ.Subscribe(monitor.Task.GID, 0)
|
||||
} else {
|
||||
monitor.setErrorStatus(errors.New("节点不可用"))
|
||||
}
|
||||
}
|
||||
|
||||
// Loop 开启监控循环
|
||||
func (monitor *Monitor) Loop() {
|
||||
defer mq.GlobalMQ.Unsubscribe(monitor.Task.GID, monitor.notifier)
|
||||
|
||||
// 首次循环立即更新
|
||||
interval := time.Duration(0)
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-monitor.notifier:
|
||||
if monitor.Update() {
|
||||
return
|
||||
}
|
||||
case <-time.After(interval):
|
||||
interval = monitor.Interval
|
||||
if monitor.Update() {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Update 更新状态，返回值表示是否退出监控
|
||||
func (monitor *Monitor) Update() bool {
|
||||
status, err := monitor.node.GetAria2Instance().Status(monitor.Task)
|
||||
|
||||
if err != nil {
|
||||
monitor.retried++
|
||||
util.Log().Warning("无法获取下载任务[%s]的状态，%s", monitor.Task.GID, err)
|
||||
|
||||
// 十次重试后认定为任务失败
|
||||
if monitor.retried > MAX_RETRY {
|
||||
util.Log().Warning("无法获取下载任务[%s]的状态，超过最大重试次数限制，%s", monitor.Task.GID, err)
|
||||
monitor.setErrorStatus(err)
|
||||
monitor.RemoveTempFolder()
|
||||
return true
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
monitor.retried = 0
|
||||
|
||||
// 磁力链下载需要跟随
|
||||
if len(status.FollowedBy) > 0 {
|
||||
util.Log().Debug("离线下载[%s]重定向至[%s]", monitor.Task.GID, status.FollowedBy[0])
|
||||
monitor.Task.GID = status.FollowedBy[0]
|
||||
monitor.Task.Save()
|
||||
return false
|
||||
}
|
||||
|
||||
// 更新任务信息
|
||||
if err := monitor.UpdateTaskInfo(status); err != nil {
|
||||
util.Log().Warning("无法更新下载任务[%s]的任务信息[%s]，", monitor.Task.GID, err)
|
||||
monitor.setErrorStatus(err)
|
||||
monitor.RemoveTempFolder()
|
||||
return true
|
||||
}
|
||||
|
||||
util.Log().Debug("离线下载[%s]更新状态[%s]", status.Gid, status.Status)
|
||||
|
||||
switch status.Status {
|
||||
case "complete":
|
||||
return monitor.Complete(status)
|
||||
case "error":
|
||||
return monitor.Error(status)
|
||||
case "active", "waiting", "paused":
|
||||
return false
|
||||
case "removed":
|
||||
monitor.Task.Status = common.Canceled
|
||||
monitor.Task.Save()
|
||||
monitor.RemoveTempFolder()
|
||||
return true
|
||||
default:
|
||||
util.Log().Warning("下载任务[%s]返回未知状态信息[%s]，", monitor.Task.GID, status.Status)
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
// UpdateTaskInfo 更新数据库中的任务信息
|
||||
func (monitor *Monitor) UpdateTaskInfo(status rpc.StatusInfo) error {
|
||||
originSize := monitor.Task.TotalSize
|
||||
|
||||
monitor.Task.GID = status.Gid
|
||||
monitor.Task.Status = common.GetStatus(status.Status)
|
||||
|
||||
// 文件大小、已下载大小
|
||||
total, err := strconv.ParseUint(status.TotalLength, 10, 64)
|
||||
if err != nil {
|
||||
total = 0
|
||||
}
|
||||
downloaded, err := strconv.ParseUint(status.CompletedLength, 10, 64)
|
||||
if err != nil {
|
||||
downloaded = 0
|
||||
}
|
||||
monitor.Task.TotalSize = total
|
||||
monitor.Task.DownloadedSize = downloaded
|
||||
monitor.Task.GID = status.Gid
|
||||
monitor.Task.Parent = status.Dir
|
||||
|
||||
// 下载速度
|
||||
speed, err := strconv.Atoi(status.DownloadSpeed)
|
||||
if err != nil {
|
||||
speed = 0
|
||||
}
|
||||
|
||||
monitor.Task.Speed = speed
|
||||
attrs, _ := json.Marshal(status)
|
||||
monitor.Task.Attrs = string(attrs)
|
||||
|
||||
if err := monitor.Task.Save(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if originSize != monitor.Task.TotalSize {
|
||||
// 文件大小更新后，对文件限制等进行校验
|
||||
if err := monitor.ValidateFile(); err != nil {
|
||||
// 验证失败时取消任务
|
||||
monitor.node.GetAria2Instance().Cancel(monitor.Task)
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ValidateFile 上传过程中校验文件大小、文件名
|
||||
func (monitor *Monitor) ValidateFile() error {
|
||||
// 找到任务创建者
|
||||
user := monitor.Task.GetOwner()
|
||||
if user == nil {
|
||||
return common.ErrUserNotFound
|
||||
}
|
||||
|
||||
// 创建文件系统
|
||||
fs, err := filesystem.NewFileSystem(user)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer fs.Recycle()
|
||||
|
||||
// 创建上下文环境
|
||||
ctx := context.WithValue(context.Background(), fsctx.FileHeaderCtx, local.FileStream{
|
||||
Size: monitor.Task.TotalSize,
|
||||
})
|
||||
|
||||
// 验证用户容量
|
||||
if err := filesystem.HookValidateCapacityWithoutIncrease(ctx, fs); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// 验证每个文件
|
||||
for _, fileInfo := range monitor.Task.StatusInfo.Files {
|
||||
if fileInfo.Selected == "true" {
|
||||
// 创建上下文环境
|
||||
fileSize, _ := strconv.ParseUint(fileInfo.Length, 10, 64)
|
||||
ctx := context.WithValue(context.Background(), fsctx.FileHeaderCtx, local.FileStream{
|
||||
Size: fileSize,
|
||||
Name: filepath.Base(fileInfo.Path),
|
||||
})
|
||||
if err := filesystem.HookValidateFile(ctx, fs); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Error 任务下载出错处理，返回是否中断监控
|
||||
func (monitor *Monitor) Error(status rpc.StatusInfo) bool {
|
||||
monitor.setErrorStatus(errors.New(status.ErrorMessage))
|
||||
|
||||
// 清理临时文件
|
||||
monitor.RemoveTempFolder()
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// RemoveTempFolder 清理下载临时目录
|
||||
func (monitor *Monitor) RemoveTempFolder() {
|
||||
monitor.node.GetAria2Instance().DeleteTempFile(monitor.Task)
|
||||
}
|
||||
|
||||
// Complete 完成下载，返回是否中断监控
|
||||
func (monitor *Monitor) Complete(status rpc.StatusInfo) bool {
|
||||
// 创建中转任务
|
||||
file := make([]string, 0, len(monitor.Task.StatusInfo.Files))
|
||||
sizes := make(map[string]uint64, len(monitor.Task.StatusInfo.Files))
|
||||
for i := 0; i < len(monitor.Task.StatusInfo.Files); i++ {
|
||||
fileInfo := monitor.Task.StatusInfo.Files[i]
|
||||
if fileInfo.Selected == "true" {
|
||||
file = append(file, fileInfo.Path)
|
||||
size, _ := strconv.ParseUint(fileInfo.Length, 10, 64)
|
||||
sizes[fileInfo.Path] = size
|
||||
}
|
||||
}
|
||||
|
||||
job, err := task.NewTransferTask(
|
||||
monitor.Task.UserID,
|
||||
file,
|
||||
monitor.Task.Dst,
|
||||
monitor.Task.Parent,
|
||||
true,
|
||||
monitor.node.ID(),
|
||||
sizes,
|
||||
)
|
||||
if err != nil {
|
||||
monitor.setErrorStatus(err)
|
||||
return true
|
||||
}
|
||||
|
||||
// 提交中转任务
|
||||
task.TaskPoll.Submit(job)
|
||||
|
||||
// 更新任务ID
|
||||
monitor.Task.TaskID = job.Model().ID
|
||||
monitor.Task.Save()
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
func (monitor *Monitor) setErrorStatus(err error) {
|
||||
monitor.Task.Status = common.Error
|
||||
monitor.Task.Error = err.Error()
|
||||
monitor.Task.Save()
|
||||
}
|
||||
326
pkg/aria2/monitor/monitor_test.go
Normal file
326
pkg/aria2/monitor/monitor_test.go
Normal file
@@ -0,0 +1,326 @@
|
||||
package monitor
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/aria2"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/aria2/common"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/aria2/rpc"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/cache"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/filesystem"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/task"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/util"
|
||||
"github.com/jinzhu/gorm"
|
||||
"github.com/stretchr/testify/assert"
|
||||
testMock "github.com/stretchr/testify/mock"
|
||||
)
|
||||
|
||||
type InstanceMock struct {
|
||||
testMock.Mock
|
||||
}
|
||||
|
||||
func (m InstanceMock) CreateTask(task *model.Download, options map[string]interface{}) error {
|
||||
args := m.Called(task, options)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m InstanceMock) Status(task *model.Download) (rpc.StatusInfo, error) {
|
||||
args := m.Called(task)
|
||||
return args.Get(0).(rpc.StatusInfo), args.Error(1)
|
||||
}
|
||||
|
||||
func (m InstanceMock) Cancel(task *model.Download) error {
|
||||
args := m.Called(task)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m InstanceMock) Select(task *model.Download, files []int) error {
|
||||
args := m.Called(task, files)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func TestNewMonitor(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
NewMonitor(&model.Download{GID: "gid"})
|
||||
_, ok := common.EventNotifier.Subscribes.Load("gid")
|
||||
asserts.True(ok)
|
||||
}
|
||||
|
||||
func TestMonitor_Loop(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
notifier := make(chan common.StatusEvent)
|
||||
MAX_RETRY = 0
|
||||
monitor := &Monitor{
|
||||
Task: &model.Download{GID: "gid"},
|
||||
Interval: time.Duration(1) * time.Second,
|
||||
notifier: notifier,
|
||||
}
|
||||
asserts.NotPanics(func() {
|
||||
monitor.Loop()
|
||||
})
|
||||
}
|
||||
|
||||
func TestMonitor_Update(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
monitor := &Monitor{
|
||||
Task: &model.Download{
|
||||
GID: "gid",
|
||||
Parent: "TestMonitor_Update",
|
||||
},
|
||||
Interval: time.Duration(1) * time.Second,
|
||||
}
|
||||
|
||||
// 无法获取状态
|
||||
{
|
||||
MAX_RETRY = 1
|
||||
testInstance := new(InstanceMock)
|
||||
testInstance.On("SlaveStatus", testMock.Anything).Return(rpc.StatusInfo{}, errors.New("error"))
|
||||
file, _ := util.CreatNestedFile("TestMonitor_Update/1")
|
||||
file.Close()
|
||||
aria2.Instance = testInstance
|
||||
asserts.False(monitor.Update())
|
||||
asserts.True(monitor.Update())
|
||||
testInstance.AssertExpectations(t)
|
||||
asserts.False(util.Exists("TestMonitor_Update"))
|
||||
}
|
||||
|
||||
// 磁力链下载重定向
|
||||
{
|
||||
testInstance := new(InstanceMock)
|
||||
testInstance.On("SlaveStatus", testMock.Anything).Return(rpc.StatusInfo{
|
||||
FollowedBy: []string{"1"},
|
||||
}, nil)
|
||||
monitor.Task.ID = 1
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
aria2.mock.ExpectCommit()
|
||||
aria2.Instance = testInstance
|
||||
asserts.False(monitor.Update())
|
||||
asserts.NoError(aria2.mock.ExpectationsWereMet())
|
||||
testInstance.AssertExpectations(t)
|
||||
asserts.EqualValues("1", monitor.Task.GID)
|
||||
}
|
||||
|
||||
// 无法更新任务信息
|
||||
{
|
||||
testInstance := new(InstanceMock)
|
||||
testInstance.On("SlaveStatus", testMock.Anything).Return(rpc.StatusInfo{}, nil)
|
||||
monitor.Task.ID = 1
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("UPDATE(.+)").WillReturnError(errors.New("error"))
|
||||
aria2.mock.ExpectRollback()
|
||||
aria2.Instance = testInstance
|
||||
asserts.True(monitor.Update())
|
||||
asserts.NoError(aria2.mock.ExpectationsWereMet())
|
||||
testInstance.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// 返回未知状态
|
||||
{
|
||||
testInstance := new(InstanceMock)
|
||||
testInstance.On("SlaveStatus", testMock.Anything).Return(rpc.StatusInfo{Status: "?"}, nil)
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
aria2.mock.ExpectCommit()
|
||||
aria2.Instance = testInstance
|
||||
asserts.True(monitor.Update())
|
||||
asserts.NoError(aria2.mock.ExpectationsWereMet())
|
||||
testInstance.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// 返回被取消状态
|
||||
{
|
||||
testInstance := new(InstanceMock)
|
||||
testInstance.On("SlaveStatus", testMock.Anything).Return(rpc.StatusInfo{Status: "removed"}, nil)
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
aria2.mock.ExpectCommit()
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
aria2.mock.ExpectCommit()
|
||||
aria2.Instance = testInstance
|
||||
asserts.True(monitor.Update())
|
||||
asserts.NoError(aria2.mock.ExpectationsWereMet())
|
||||
testInstance.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// 返回活跃状态
|
||||
{
|
||||
testInstance := new(InstanceMock)
|
||||
testInstance.On("SlaveStatus", testMock.Anything).Return(rpc.StatusInfo{Status: "active"}, nil)
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
aria2.mock.ExpectCommit()
|
||||
aria2.Instance = testInstance
|
||||
asserts.False(monitor.Update())
|
||||
asserts.NoError(aria2.mock.ExpectationsWereMet())
|
||||
testInstance.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// 返回错误状态
|
||||
{
|
||||
testInstance := new(InstanceMock)
|
||||
testInstance.On("SlaveStatus", testMock.Anything).Return(rpc.StatusInfo{Status: "error"}, nil)
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
aria2.mock.ExpectCommit()
|
||||
aria2.Instance = testInstance
|
||||
asserts.True(monitor.Update())
|
||||
asserts.NoError(aria2.mock.ExpectationsWereMet())
|
||||
testInstance.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// 返回完成
|
||||
{
|
||||
testInstance := new(InstanceMock)
|
||||
testInstance.On("SlaveStatus", testMock.Anything).Return(rpc.StatusInfo{Status: "complete"}, nil)
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
aria2.mock.ExpectCommit()
|
||||
aria2.Instance = testInstance
|
||||
asserts.True(monitor.Update())
|
||||
asserts.NoError(aria2.mock.ExpectationsWereMet())
|
||||
testInstance.AssertExpectations(t)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMonitor_UpdateTaskInfo(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
monitor := &Monitor{
|
||||
Task: &model.Download{
|
||||
Model: gorm.Model{ID: 1},
|
||||
GID: "gid",
|
||||
Parent: "TestMonitor_UpdateTaskInfo",
|
||||
},
|
||||
}
|
||||
|
||||
// 失败
|
||||
{
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("UPDATE(.+)").WillReturnError(errors.New("error"))
|
||||
aria2.mock.ExpectRollback()
|
||||
err := monitor.UpdateTaskInfo(rpc.StatusInfo{})
|
||||
asserts.NoError(aria2.mock.ExpectationsWereMet())
|
||||
asserts.Error(err)
|
||||
}
|
||||
|
||||
// 更新成功，无需校验
|
||||
{
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
aria2.mock.ExpectCommit()
|
||||
err := monitor.UpdateTaskInfo(rpc.StatusInfo{})
|
||||
asserts.NoError(aria2.mock.ExpectationsWereMet())
|
||||
asserts.NoError(err)
|
||||
}
|
||||
|
||||
// 更新成功，大小改变，需要校验，校验失败
|
||||
{
|
||||
testInstance := new(InstanceMock)
|
||||
testInstance.On("SlaveCancel", testMock.Anything).Return(nil)
|
||||
aria2.Instance = testInstance
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
aria2.mock.ExpectCommit()
|
||||
err := monitor.UpdateTaskInfo(rpc.StatusInfo{TotalLength: "1"})
|
||||
asserts.NoError(aria2.mock.ExpectationsWereMet())
|
||||
asserts.Error(err)
|
||||
testInstance.AssertExpectations(t)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMonitor_ValidateFile(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
monitor := &Monitor{
|
||||
Task: &model.Download{
|
||||
Model: gorm.Model{ID: 1},
|
||||
GID: "gid",
|
||||
Parent: "TestMonitor_ValidateFile",
|
||||
},
|
||||
}
|
||||
|
||||
// 无法创建文件系统
|
||||
{
|
||||
monitor.Task.User = &model.User{
|
||||
Policy: model.Policy{
|
||||
Type: "unknown",
|
||||
},
|
||||
}
|
||||
asserts.Error(monitor.ValidateFile())
|
||||
}
|
||||
|
||||
// 文件大小超出容量配额
|
||||
{
|
||||
cache.Set("pack_size_0", uint64(0), 0)
|
||||
monitor.Task.TotalSize = 11
|
||||
monitor.Task.User = &model.User{
|
||||
Policy: model.Policy{
|
||||
Type: "mock",
|
||||
},
|
||||
Group: model.Group{
|
||||
MaxStorage: 10,
|
||||
},
|
||||
}
|
||||
asserts.Equal(filesystem.ErrInsufficientCapacity, monitor.ValidateFile())
|
||||
}
|
||||
|
||||
// 单文件大小超出容量配额
|
||||
{
|
||||
cache.Set("pack_size_0", uint64(0), 0)
|
||||
monitor.Task.TotalSize = 10
|
||||
monitor.Task.StatusInfo.Files = []rpc.FileInfo{
|
||||
{
|
||||
Selected: "true",
|
||||
Length: "6",
|
||||
},
|
||||
}
|
||||
monitor.Task.User = &model.User{
|
||||
Policy: model.Policy{
|
||||
Type: "mock",
|
||||
MaxSize: 5,
|
||||
},
|
||||
Group: model.Group{
|
||||
MaxStorage: 10,
|
||||
},
|
||||
}
|
||||
asserts.Equal(filesystem.ErrFileSizeTooBig, monitor.ValidateFile())
|
||||
}
|
||||
}
|
||||
|
||||
func TestMonitor_Complete(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
monitor := &Monitor{
|
||||
Task: &model.Download{
|
||||
Model: gorm.Model{ID: 1},
|
||||
GID: "gid",
|
||||
Parent: "TestMonitor_Complete",
|
||||
StatusInfo: rpc.StatusInfo{
|
||||
Files: []rpc.FileInfo{
|
||||
{
|
||||
Selected: "true",
|
||||
Path: "TestMonitor_Complete",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
cache.Set("setting_max_worker_num", "1", 0)
|
||||
aria2.mock.ExpectQuery("SELECT(.+)tasks").WillReturnRows(sqlmock.NewRows([]string{"id"}))
|
||||
task.Init()
|
||||
aria2.mock.ExpectQuery("SELECT(.+)users").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
aria2.mock.ExpectQuery("SELECT(.+)policies").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("INSERT(.+)tasks").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
aria2.mock.ExpectCommit()
|
||||
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("UPDATE(.+)downloads").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
aria2.mock.ExpectCommit()
|
||||
asserts.True(monitor.Complete(rpc.StatusInfo{}))
|
||||
asserts.NoError(aria2.mock.ExpectationsWereMet())
|
||||
}
|
||||
Reference in New Issue
Block a user