github.com/cloudreve/Cloudreve/v3@v3.0.0-20240224133659-3edb00a6484c/pkg/aria2/monitor/monitor_test.go

package monitor

import (
	"database/sql"
	"errors"
	"testing"

	"github.com/DATA-DOG/go-sqlmock"
	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/filesystem"
	"github.com/cloudreve/Cloudreve/v3/pkg/mocks"
	"github.com/cloudreve/Cloudreve/v3/pkg/mq"
	"github.com/jinzhu/gorm"
	"github.com/stretchr/testify/assert"
	testMock "github.com/stretchr/testify/mock"
)

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 TestNewMonitor(t *testing.T) {
	a := assert.New(t)
	mockMQ := mq.NewMQ()

	// node not available
	{
		mockPool := &mocks.NodePoolMock{}
		mockPool.On("GetNodeByID", uint(1)).Return(nil)
		mock.ExpectBegin()
		mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
		mock.ExpectCommit()

		task := &model.Download{
			Model: gorm.Model{ID: 1},
		}
		NewMonitor(task, mockPool, mockMQ)
		mockPool.AssertExpectations(t)
		a.NoError(mock.ExpectationsWereMet())
		a.NotEmpty(task.Error)
	}

	// success
	{
		mockNode := &mocks.NodeMock{}
		mockNode.On("GetAria2Instance").Return(&common.DummyAria2{})
		mockPool := &mocks.NodePoolMock{}
		mockPool.On("GetNodeByID", uint(1)).Return(mockNode)

		task := &model.Download{
			Model: gorm.Model{ID: 1},
		}
		NewMonitor(task, mockPool, mockMQ)
		mockNode.AssertExpectations(t)
		mockPool.AssertExpectations(t)
	}

}

func TestMonitor_Loop(t *testing.T) {
	a := assert.New(t)
	mockMQ := mq.NewMQ()
	mockNode := &mocks.NodeMock{}
	mockNode.On("GetAria2Instance").Return(&common.DummyAria2{})
	m := &Monitor{
		retried:  MAX_RETRY,
		node:     mockNode,
		Task:     &model.Download{Model: gorm.Model{ID: 1}},
		notifier: mockMQ.Subscribe("test", 1),
	}

	// into interval loop
	{
		mock.ExpectBegin()
		mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
		mock.ExpectCommit()
		m.Loop(mockMQ)
		a.NoError(mock.ExpectationsWereMet())
		a.NotEmpty(m.Task.Error)
	}

	// into notifier loop
	{
		m.Task.Error = ""
		mockMQ.Publish("test", mq.Message{})
		mock.ExpectBegin()
		mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
		mock.ExpectCommit()
		m.Loop(mockMQ)
		a.NoError(mock.ExpectationsWereMet())
		a.NotEmpty(m.Task.Error)
	}
}

func TestMonitor_UpdateFailedAfterRetry(t *testing.T) {
	a := assert.New(t)
	mockNode := &mocks.NodeMock{}
	mockNode.On("GetAria2Instance").Return(&common.DummyAria2{})
	m := &Monitor{
		node: mockNode,
		Task: &model.Download{Model: gorm.Model{ID: 1}},
	}
	mock.ExpectBegin()
	mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
	mock.ExpectCommit()

	for i := 0; i < MAX_RETRY; i++ {
		a.False(m.Update())
	}

	mockNode.AssertExpectations(t)
	a.True(m.Update())
	a.NoError(mock.ExpectationsWereMet())
	a.NotEmpty(m.Task.Error)
}

func TestMonitor_UpdateMagentoFollow(t *testing.T) {
	a := assert.New(t)
	mockAria2 := &mocks.Aria2Mock{}
	mockAria2.On("Status", testMock.Anything).Return(rpc.StatusInfo{
		FollowedBy: []string{"next"},
	}, nil)
	mockNode := &mocks.NodeMock{}
	mockNode.On("GetAria2Instance").Return(mockAria2)
	m := &Monitor{
		node: mockNode,
		Task: &model.Download{Model: gorm.Model{ID: 1}},
	}
	mock.ExpectBegin()
	mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
	mock.ExpectCommit()

	a.False(m.Update())
	a.NoError(mock.ExpectationsWereMet())
	a.Equal("next", m.Task.GID)
	mockAria2.AssertExpectations(t)
}

func TestMonitor_UpdateFailedToUpdateInfo(t *testing.T) {
	a := assert.New(t)
	mockAria2 := &mocks.Aria2Mock{}
	mockAria2.On("Status", testMock.Anything).Return(rpc.StatusInfo{}, nil)
	mockAria2.On("DeleteTempFile", testMock.Anything).Return(nil)
	mockNode := &mocks.NodeMock{}
	mockNode.On("GetAria2Instance").Return(mockAria2)
	m := &Monitor{
		node: mockNode,
		Task: &model.Download{Model: gorm.Model{ID: 1}},
	}
	mock.ExpectBegin()
	mock.ExpectExec("UPDATE(.+)").WillReturnError(errors.New("error"))
	mock.ExpectRollback()
	mock.ExpectBegin()
	mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
	mock.ExpectCommit()

	a.True(m.Update())
	a.NoError(mock.ExpectationsWereMet())
	mockAria2.AssertExpectations(t)
	mockNode.AssertExpectations(t)
	a.NotEmpty(m.Task.Error)
}

func TestMonitor_UpdateCompleted(t *testing.T) {
	a := assert.New(t)
	mockAria2 := &mocks.Aria2Mock{}
	mockAria2.On("Status", testMock.Anything).Return(rpc.StatusInfo{
		Status: "complete",
	}, nil)
	mockAria2.On("DeleteTempFile", testMock.Anything).Return(nil)
	mockNode := &mocks.NodeMock{}
	mockNode.On("GetAria2Instance").Return(mockAria2)
	mockNode.On("ID").Return(uint(1))
	m := &Monitor{
		node: mockNode,
		Task: &model.Download{Model: gorm.Model{ID: 1}},
	}
	mock.ExpectBegin()
	mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
	mock.ExpectCommit()
	mock.ExpectQuery("SELECT(.+)users(.+)").WillReturnError(errors.New("error"))
	mock.ExpectBegin()
	mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
	mock.ExpectCommit()

	a.True(m.Update())
	a.NoError(mock.ExpectationsWereMet())
	mockAria2.AssertExpectations(t)
	mockNode.AssertExpectations(t)
	a.NotEmpty(m.Task.Error)
}

func TestMonitor_UpdateError(t *testing.T) {
	a := assert.New(t)
	mockAria2 := &mocks.Aria2Mock{}
	mockAria2.On("Status", testMock.Anything).Return(rpc.StatusInfo{
		Status:       "error",
		ErrorMessage: "error",
	}, nil)
	mockAria2.On("DeleteTempFile", testMock.Anything).Return(nil)
	mockNode := &mocks.NodeMock{}
	mockNode.On("GetAria2Instance").Return(mockAria2)
	m := &Monitor{
		node: mockNode,
		Task: &model.Download{Model: gorm.Model{ID: 1}},
	}
	mock.ExpectBegin()
	mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
	mock.ExpectCommit()
	mock.ExpectBegin()
	mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
	mock.ExpectCommit()

	a.True(m.Update())
	a.NoError(mock.ExpectationsWereMet())
	mockAria2.AssertExpectations(t)
	mockNode.AssertExpectations(t)
	a.NotEmpty(m.Task.Error)
}

func TestMonitor_UpdateActive(t *testing.T) {
	a := assert.New(t)
	mockAria2 := &mocks.Aria2Mock{}
	mockAria2.On("Status", testMock.Anything).Return(rpc.StatusInfo{
		Status: "active",
	}, nil)
	mockNode := &mocks.NodeMock{}
	mockNode.On("GetAria2Instance").Return(mockAria2)
	m := &Monitor{
		node: mockNode,
		Task: &model.Download{Model: gorm.Model{ID: 1}},
	}
	mock.ExpectBegin()
	mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
	mock.ExpectCommit()

	a.False(m.Update())
	a.NoError(mock.ExpectationsWereMet())
	mockAria2.AssertExpectations(t)
	mockNode.AssertExpectations(t)
}

func TestMonitor_UpdateRemoved(t *testing.T) {
	a := assert.New(t)
	mockAria2 := &mocks.Aria2Mock{}
	mockAria2.On("Status", testMock.Anything).Return(rpc.StatusInfo{
		Status: "removed",
	}, nil)
	mockAria2.On("DeleteTempFile", testMock.Anything).Return(nil)
	mockNode := &mocks.NodeMock{}
	mockNode.On("GetAria2Instance").Return(mockAria2)
	m := &Monitor{
		node: mockNode,
		Task: &model.Download{Model: gorm.Model{ID: 1}},
	}
	mock.ExpectBegin()
	mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
	mock.ExpectCommit()

	mock.ExpectBegin()
	mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
	mock.ExpectCommit()

	a.True(m.Update())
	a.Equal(common.Canceled, m.Task.Status)
	a.NoError(mock.ExpectationsWereMet())
	mockAria2.AssertExpectations(t)
	mockNode.AssertExpectations(t)
}

func TestMonitor_UpdateUnknown(t *testing.T) {
	a := assert.New(t)
	mockAria2 := &mocks.Aria2Mock{}
	mockAria2.On("Status", testMock.Anything).Return(rpc.StatusInfo{
		Status: "unknown",
	}, nil)
	mockNode := &mocks.NodeMock{}
	mockNode.On("GetAria2Instance").Return(mockAria2)
	m := &Monitor{
		node: mockNode,
		Task: &model.Download{Model: gorm.Model{ID: 1}},
	}
	mock.ExpectBegin()
	mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
	mock.ExpectCommit()

	a.True(m.Update())
	a.NoError(mock.ExpectationsWereMet())
	mockAria2.AssertExpectations(t)
	mockNode.AssertExpectations(t)
}

func TestMonitor_UpdateTaskInfoValidateFailed(t *testing.T) {
	a := assert.New(t)
	status := rpc.StatusInfo{
		Status:          "completed",
		TotalLength:     "100",
		CompletedLength: "50",
		DownloadSpeed:   "20",
	}
	mockNode := &mocks.NodeMock{}
	mockNode.On("GetAria2Instance").Return(&common.DummyAria2{})
	m := &Monitor{
		node: mockNode,
		Task: &model.Download{Model: gorm.Model{ID: 1}},
	}

	mock.ExpectBegin()
	mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
	mock.ExpectCommit()

	err := m.UpdateTaskInfo(status)
	a.Error(err)
	a.NoError(mock.ExpectationsWereMet())
	mockNode.AssertExpectations(t)
}

func TestMonitor_ValidateFile(t *testing.T) {
	a := assert.New(t)
	m := &Monitor{
		Task: &model.Download{
			Model:     gorm.Model{ID: 1},
			TotalSize: 100,
		},
	}

	// failed to create filesystem
	{
		m.Task.User = &model.User{
			Policy: model.Policy{
				Type: "random",
			},
		}
		a.Equal(filesystem.ErrUnknownPolicyType, m.ValidateFile())
	}

	// User capacity not enough
	{
		m.Task.User = &model.User{
			Group: model.Group{
				MaxStorage: 99,
			},
			Policy: model.Policy{
				Type: "local",
			},
		}
		a.Equal(filesystem.ErrInsufficientCapacity, m.ValidateFile())
	}

	// single file too big
	{
		m.Task.StatusInfo.Files = []rpc.FileInfo{
			{
				Length:   "100",
				Selected: "true",
			},
		}
		m.Task.User = &model.User{
			Group: model.Group{
				MaxStorage: 100,
			},
			Policy: model.Policy{
				Type:    "local",
				MaxSize: 99,
			},
		}
		a.Equal(filesystem.ErrFileSizeTooBig, m.ValidateFile())
	}

	// all pass
	{
		m.Task.StatusInfo.Files = []rpc.FileInfo{
			{
				Length:   "100",
				Selected: "true",
			},
		}
		m.Task.User = &model.User{
			Group: model.Group{
				MaxStorage: 100,
			},
			Policy: model.Policy{
				Type:    "local",
				MaxSize: 100,
			},
		}
		a.NoError(m.ValidateFile())
	}
}

func TestMonitor_Complete(t *testing.T) {
	a := assert.New(t)
	mockNode := &mocks.NodeMock{}
	mockNode.On("ID").Return(uint(1))
	mockPool := &mocks.TaskPoolMock{}
	mockPool.On("Submit", testMock.Anything)
	m := &Monitor{
		node: mockNode,
		Task: &model.Download{
			Model:     gorm.Model{ID: 1},
			TotalSize: 100,
			UserID:    9414,
		},
	}
	m.Task.StatusInfo.Files = []rpc.FileInfo{
		{
			Length:   "100",
			Selected: "true",
		},
	}

	mock.ExpectQuery("SELECT(.+)users").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(9414))

	mock.ExpectBegin()
	mock.ExpectExec("INSERT(.+)tasks").WillReturnResult(sqlmock.NewResult(1, 1))
	mock.ExpectCommit()

	mock.ExpectBegin()
	mock.ExpectExec("UPDATE(.+)downloads").WillReturnResult(sqlmock.NewResult(1, 1))
	mock.ExpectCommit()

	mock.ExpectQuery("SELECT(.+)tasks").WillReturnRows(sqlmock.NewRows([]string{"id", "type", "status"}).AddRow(1, 2, 4))
	mock.ExpectQuery("SELECT(.+)users").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(9414))
	mock.ExpectBegin()
	mock.ExpectExec("INSERT(.+)tasks").WillReturnResult(sqlmock.NewResult(2, 1))
	mock.ExpectCommit()

	a.False(m.Complete(mockPool))
	m.Task.StatusInfo.Status = "complete"
	a.True(m.Complete(mockPool))
	a.NoError(mock.ExpectationsWereMet())
	mockNode.AssertExpectations(t)
	mockPool.AssertExpectations(t)
}