From 0cfa61e264f0773da48ef214600f43d31b52d304 Mon Sep 17 00:00:00 2001
From: HFO4 <912394456@qq.com>
Date: Sun, 6 Dec 2020 16:50:08 +0800
Subject: [PATCH] Test: user storage calibration script

---
 models/scripts/invoker_test.go | 49 ++++++++++++++++++++++++++++++++++
 models/scripts/storage_test.go | 38 ++++++++++++++++++++++++++
 2 files changed, 87 insertions(+)
 create mode 100644 models/scripts/invoker_test.go
 create mode 100644 models/scripts/storage_test.go

diff --git a/models/scripts/invoker_test.go b/models/scripts/invoker_test.go
new file mode 100644
index 0000000..0ca324b
--- /dev/null
+++ b/models/scripts/invoker_test.go
@@ -0,0 +1,49 @@
+package scripts
+
+import (
+	"context"
+	"database/sql"
+	"github.com/DATA-DOG/go-sqlmock"
+	model "github.com/cloudreve/Cloudreve/v3/models"
+	"github.com/jinzhu/gorm"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+var mock sqlmock.Sqlmock
+var mockDB *gorm.DB
+
+type TestScript int
+
+func (script TestScript) Run(ctx context.Context) {
+
+}
+
+// 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)
+	mockDB = model.DB
+	defer db.Close()
+	m.Run()
+}
+
+func TestRunDBScript(t *testing.T) {
+	asserts := assert.New(t)
+	register("test", TestScript(0))
+
+	// 不存在
+	{
+		asserts.Error(RunDBScript("else", context.Background()))
+	}
+
+	// 存在
+	{
+		asserts.NoError(RunDBScript("test", context.Background()))
+	}
+}
diff --git a/models/scripts/storage_test.go b/models/scripts/storage_test.go
new file mode 100644
index 0000000..7287724
--- /dev/null
+++ b/models/scripts/storage_test.go
@@ -0,0 +1,38 @@
+package scripts
+
+import (
+	"context"
+	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestUserStorageCalibration_Run(t *testing.T) {
+	asserts := assert.New(t)
+	script := UserStorageCalibration(0)
+
+	// 容量异常
+	{
+		mock.ExpectQuery("SELECT(.+)users(.+)").
+			WillReturnRows(sqlmock.NewRows([]string{"id", "email", "storage"}).AddRow(1, "a@a.com", 10))
+		mock.ExpectQuery("SELECT(.+)files(.+)").
+			WithArgs(1).
+			WillReturnRows(sqlmock.NewRows([]string{"total"}).AddRow(11))
+		mock.ExpectBegin()
+		mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
+		mock.ExpectCommit()
+		script.Run(context.Background())
+		asserts.NoError(mock.ExpectationsWereMet())
+	}
+
+	// 容量正常
+	{
+		mock.ExpectQuery("SELECT(.+)users(.+)").
+			WillReturnRows(sqlmock.NewRows([]string{"id", "email", "storage"}).AddRow(1, "a@a.com", 10))
+		mock.ExpectQuery("SELECT(.+)files(.+)").
+			WithArgs(1).
+			WillReturnRows(sqlmock.NewRows([]string{"total"}).AddRow(10))
+		script.Run(context.Background())
+		asserts.NoError(mock.ExpectationsWereMet())
+	}
+}