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  }