github.com/cloudreve/Cloudreve/v3@v3.0.0-20240224133659-3edb00a6484c/pkg/task/job_test.go (about)

     1  package task
     2  
     3  import (
     4  	"errors"
     5  	"testing"
     6  
     7  	"github.com/DATA-DOG/go-sqlmock"
     8  	model "github.com/cloudreve/Cloudreve/v3/models"
     9  	"github.com/stretchr/testify/assert"
    10  	testMock "github.com/stretchr/testify/mock"
    11  )
    12  
    13  func TestRecord(t *testing.T) {
    14  	asserts := assert.New(t)
    15  	job := &TransferTask{
    16  		User: &model.User{Policy: model.Policy{Type: "unknown"}},
    17  	}
    18  	mock.ExpectBegin()
    19  	mock.ExpectExec("INSERT(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
    20  	mock.ExpectCommit()
    21  	_, err := Record(job)
    22  	asserts.NoError(err)
    23  }
    24  
    25  type taskPoolMock struct {
    26  	testMock.Mock
    27  }
    28  
    29  func (t taskPoolMock) Add(num int) {
    30  	t.Called(num)
    31  }
    32  
    33  func (t taskPoolMock) Submit(job Job) {
    34  	t.Called(job)
    35  }
    36  
    37  func TestResume(t *testing.T) {
    38  	asserts := assert.New(t)
    39  	mockPool := taskPoolMock{}
    40  
    41  	// 没有任务
    42  	{
    43  		mock.ExpectQuery("SELECT(.+)").WithArgs(Queued, Processing).WillReturnRows(sqlmock.NewRows([]string{"type"}))
    44  		Resume(mockPool)
    45  		asserts.NoError(mock.ExpectationsWereMet())
    46  	}
    47  
    48  	// 有任务, 类型未知
    49  	{
    50  		mock.ExpectQuery("SELECT(.+)").WithArgs(Queued, Processing).WillReturnRows(sqlmock.NewRows([]string{"type"}).AddRow(233))
    51  		Resume(mockPool)
    52  		asserts.NoError(mock.ExpectationsWereMet())
    53  	}
    54  
    55  	// 有任务
    56  	{
    57  		mockPool.On("Submit", testMock.Anything)
    58  		mock.ExpectQuery("SELECT(.+)").WithArgs(Queued, Processing).WillReturnRows(sqlmock.NewRows([]string{"type", "props"}).AddRow(CompressTaskType, "{}"))
    59  		mock.ExpectQuery("SELECT(.+)users").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
    60  		mock.ExpectQuery("SELECT(.+)policies").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
    61  		Resume(mockPool)
    62  		asserts.NoError(mock.ExpectationsWereMet())
    63  		mockPool.AssertExpectations(t)
    64  	}
    65  }
    66  
    67  func TestGetJobFromModel(t *testing.T) {
    68  	asserts := assert.New(t)
    69  
    70  	// CompressTaskType
    71  	{
    72  		task := &model.Task{
    73  			Status: 0,
    74  			Type:   CompressTaskType,
    75  		}
    76  		mock.ExpectQuery("SELECT(.+)users(.+)").WillReturnError(errors.New("error"))
    77  		job, err := GetJobFromModel(task)
    78  		asserts.NoError(mock.ExpectationsWereMet())
    79  		asserts.Nil(job)
    80  		asserts.Error(err)
    81  	}
    82  	// DecompressTaskType
    83  	{
    84  		task := &model.Task{
    85  			Status: 0,
    86  			Type:   DecompressTaskType,
    87  		}
    88  		mock.ExpectQuery("SELECT(.+)users(.+)").WillReturnError(errors.New("error"))
    89  		job, err := GetJobFromModel(task)
    90  		asserts.NoError(mock.ExpectationsWereMet())
    91  		asserts.Nil(job)
    92  		asserts.Error(err)
    93  	}
    94  	// TransferTaskType
    95  	{
    96  		task := &model.Task{
    97  			Status: 0,
    98  			Type:   TransferTaskType,
    99  		}
   100  		mock.ExpectQuery("SELECT(.+)users(.+)").WillReturnError(errors.New("error"))
   101  		job, err := GetJobFromModel(task)
   102  		asserts.NoError(mock.ExpectationsWereMet())
   103  		asserts.Nil(job)
   104  		asserts.Error(err)
   105  	}
   106  	// RecycleTaskType
   107  	{
   108  		task := &model.Task{
   109  			Status: 0,
   110  			Type:   RecycleTaskType,
   111  		}
   112  		mock.ExpectQuery("SELECT(.+)users(.+)").WillReturnError(errors.New("error"))
   113  		job, err := GetJobFromModel(task)
   114  		asserts.NoError(mock.ExpectationsWereMet())
   115  		asserts.Nil(job)
   116  		asserts.Error(err)
   117  	}
   118  }