From 8dd411f5d4e8baa72c6669f140d8955724db2bd4 Mon Sep 17 00:00:00 2001
From: HFO4 <912394456@qq.com>
Date: Sat, 16 Nov 2019 16:58:10 +0800
Subject: [PATCH] Feat: memory database for testing

---
 models/init.go         | 22 +++++++++++++++-------
 routers/router_test.go |  5 +++--
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/models/init.go b/models/init.go
index a25d3b9..fb8fe9f 100644
--- a/models/init.go
+++ b/models/init.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"github.com/HFO4/cloudreve/pkg/conf"
 	"github.com/HFO4/cloudreve/pkg/util"
+	"github.com/gin-gonic/gin"
 	"github.com/jinzhu/gorm"
 	"time"
 
@@ -22,14 +23,21 @@ func Init() {
 		db  *gorm.DB
 		err error
 	)
-	if conf.DatabaseConfig.Type == "UNSET" {
-		db, err = gorm.Open("sqlite3", "cloudreve.db")
+
+	if gin.Mode() == gin.TestMode {
+		// 测试模式下，使用内存数据库
+		db, err = gorm.Open("sqlite3", ":memory:")
 	} else {
-		db, err = gorm.Open(conf.DatabaseConfig.Type, fmt.Sprintf("%s:%s@(%s)/%s?charset=utf8&parseTime=True&loc=Local",
-			conf.DatabaseConfig.User,
-			conf.DatabaseConfig.Password,
-			conf.DatabaseConfig.Host,
-			conf.DatabaseConfig.Name))
+		if conf.DatabaseConfig.Type == "UNSET" {
+			// 未指定数据库时，使用Sqlite
+			db, err = gorm.Open("sqlite3", "cloudreve.db")
+		} else {
+			db, err = gorm.Open(conf.DatabaseConfig.Type, fmt.Sprintf("%s:%s@(%s)/%s?charset=utf8&parseTime=True&loc=Local",
+				conf.DatabaseConfig.User,
+				conf.DatabaseConfig.Password,
+				conf.DatabaseConfig.Host,
+				conf.DatabaseConfig.Name))
+		}
 	}
 
 	// 处理表前缀
diff --git a/routers/router_test.go b/routers/router_test.go
index 3b70ee5..b45ccda 100644
--- a/routers/router_test.go
+++ b/routers/router_test.go
@@ -22,6 +22,9 @@ var mock sqlmock.Sqlmock
 
 // TestMain 初始化数据库Mock
 func TestMain(m *testing.M) {
+	// 设置gin为测试模式
+	gin.SetMode(gin.TestMode)
+
 	var db *sql.DB
 	var err error
 	db, mock, err = sqlmock.New()
@@ -31,8 +34,6 @@ func TestMain(m *testing.M) {
 	model.DB, _ = gorm.Open("mysql", db)
 	defer db.Close()
 
-	// 设置gin为测试模式
-	gin.SetMode(gin.TestMode)
 	m.Run()
 }