github.com/cloudreve/Cloudreve/v3@v3.0.0-20240224133659-3edb00a6484c/models/scripts/storage_test.go (about) 1 package scripts 2 3 import ( 4 "context" 5 "database/sql" 6 "github.com/DATA-DOG/go-sqlmock" 7 model "github.com/cloudreve/Cloudreve/v3/models" 8 "github.com/jinzhu/gorm" 9 "github.com/stretchr/testify/assert" 10 "testing" 11 ) 12 13 var mock sqlmock.Sqlmock 14 var mockDB *gorm.DB 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 mockDB = model.DB 26 defer db.Close() 27 m.Run() 28 } 29 30 func TestUserStorageCalibration_Run(t *testing.T) { 31 asserts := assert.New(t) 32 script := UserStorageCalibration(0) 33 34 // 容量异常 35 { 36 mock.ExpectQuery("SELECT(.+)users(.+)"). 37 WillReturnRows(sqlmock.NewRows([]string{"id", "email", "storage"}).AddRow(1, "a@a.com", 10)) 38 mock.ExpectQuery("SELECT(.+)files(.+)"). 39 WithArgs(1). 40 WillReturnRows(sqlmock.NewRows([]string{"total"}).AddRow(11)) 41 mock.ExpectBegin() 42 mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1)) 43 mock.ExpectCommit() 44 script.Run(context.Background()) 45 asserts.NoError(mock.ExpectationsWereMet()) 46 } 47 48 // 容量正常 49 { 50 mock.ExpectQuery("SELECT(.+)users(.+)"). 51 WillReturnRows(sqlmock.NewRows([]string{"id", "email", "storage"}).AddRow(1, "a@a.com", 10)) 52 mock.ExpectQuery("SELECT(.+)files(.+)"). 53 WithArgs(1). 54 WillReturnRows(sqlmock.NewRows([]string{"total"}).AddRow(10)) 55 mock.ExpectBegin() 56 mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1)) 57 mock.ExpectCommit() 58 script.Run(context.Background()) 59 asserts.NoError(mock.ExpectationsWereMet()) 60 } 61 }