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

     1  package task
     2  
     3  import (
     4  	"database/sql"
     5  	"testing"
     6  
     7  	"github.com/DATA-DOG/go-sqlmock"
     8  	model "github.com/cloudreve/Cloudreve/v3/models"
     9  	"github.com/cloudreve/Cloudreve/v3/pkg/cache"
    10  	"github.com/jinzhu/gorm"
    11  	"github.com/stretchr/testify/assert"
    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 TestInit(t *testing.T) {
    30  	asserts := assert.New(t)
    31  	cache.Set("setting_max_worker_num", "10", 0)
    32  	mock.ExpectQuery("SELECT(.+)").WithArgs(Queued, Processing).WillReturnRows(sqlmock.NewRows([]string{"type"}).AddRow(-1))
    33  	Init()
    34  	asserts.NoError(mock.ExpectationsWereMet())
    35  	asserts.Len(TaskPoll.(*AsyncPool).idleWorker, 10)
    36  }
    37  
    38  func TestPool_Submit(t *testing.T) {
    39  	asserts := assert.New(t)
    40  	pool := &AsyncPool{
    41  		idleWorker: make(chan int, 1),
    42  	}
    43  	pool.Add(1)
    44  	job := &MockJob{
    45  		DoFunc: func() {
    46  
    47  		},
    48  	}
    49  	asserts.NotPanics(func() {
    50  		pool.Submit(job)
    51  	})
    52  }