cloudreve/routers/router_test.go

package routers

import (
	"bytes"
	"database/sql"
	"encoding/json"
	"errors"
	"github.com/DATA-DOG/go-sqlmock"
	"github.com/HFO4/cloudreve/middleware"
	"github.com/HFO4/cloudreve/models"
	"github.com/HFO4/cloudreve/pkg/serializer"
	"github.com/gin-gonic/gin"
	"github.com/jinzhu/gorm"
	"github.com/mojocn/base64Captcha"
	"github.com/stretchr/testify/assert"
	"net/http"
	"net/http/httptest"
	"testing"
)

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()
	if err != nil {
		panic("An error was not expected when opening a stub database connection")
	}
	model.DB, _ = gorm.Open("mysql", db)
	defer db.Close()

	m.Run()
}

func TestPing(t *testing.T) {
	asserts := assert.New(t)
	router := InitRouter()

	w := httptest.NewRecorder()
	req, _ := http.NewRequest("GET", "/Api/V3/Ping", nil)
	router.ServeHTTP(w, req)

	assert.Equal(t, 200, w.Code)
	asserts.Contains(w.Body.String(), "Pong")
}

func TestCaptcha(t *testing.T) {
	asserts := assert.New(t)
	router := InitRouter()
	w := httptest.NewRecorder()

	req, _ := http.NewRequest(
		"GET",
		"/Api/V3/Captcha",
		nil,
	)

	router.ServeHTTP(w, req)

	asserts.Equal(200, w.Code)
	asserts.Contains(w.Body.String(), "base64")
}

func TestUserSession(t *testing.T) {
	asserts := assert.New(t)
	router := InitRouter()
	w := httptest.NewRecorder()

	// 创建测试用验证码
	var configD = base64Captcha.ConfigDigit{
		Height:     80,
		Width:      240,
		MaxSkew:    0.7,
		DotCount:   80,
		CaptchaLen: 1,
	}
	idKeyD, _ := base64Captcha.GenerateCaptcha("", configD)
	middleware.ContextMock = map[string]interface{}{
		"captchaID": idKeyD,
	}

	testCases := []struct {
		settingRows *sqlmock.Rows
		userRows    *sqlmock.Rows
		policyRows  *sqlmock.Rows
		reqBody     string
		expected    interface{}
	}{
		// 登录信息正确，不需要验证码
		{
			settingRows: sqlmock.NewRows([]string{"name", "value", "type"}).
				AddRow("login_captcha", "0", "login"),
			userRows: sqlmock.NewRows([]string{"email", "nick", "password", "options"}).
				AddRow("admin@cloudreve.org", "admin", "CKLmDKa1C9SD64vU:76adadd4fd4bad86959155f6f7bc8993c94e7adf", "{}"), policyRows: sqlmock.NewRows([]string{"name", "type", "options"}).
				AddRow("默认上传策略", "local", "{\"op_name\":\"123\"}"),
			reqBody: `{"userName":"admin@cloudreve.org","captchaCode":"captchaCode","Password":"admin"}`,
			expected: serializer.BuildUserResponse(model.User{
				Email: "admin@cloudreve.org",
				Nick:  "admin",
			}),
		},
		// 登录信息正确，需要验证码,验证码错误
		{
			settingRows: sqlmock.NewRows([]string{"name", "value", "type"}).
				AddRow("login_captcha", "1", "login"),
			userRows: sqlmock.NewRows([]string{"email", "nick", "password", "options"}).
				AddRow("admin@cloudreve.org", "admin", "CKLmDKa1C9SD64vU:76adadd4fd4bad86959155f6f7bc8993c94e7adf", "{}"),
			policyRows: sqlmock.NewRows([]string{"name", "type", "options"}).
				AddRow("默认上传策略", "local", "{\"op_name\":\"123\"}"),
			reqBody:  `{"userName":"admin@cloudreve.org","captchaCode":"captchaCode","Password":"admin"}`,
			expected: serializer.ParamErr("验证码错误", nil),
		},
		// 邮箱正确密码错误
		{
			settingRows: sqlmock.NewRows([]string{"name", "value", "type"}).
				AddRow("login_captcha", "0", "login"),
			userRows: sqlmock.NewRows([]string{"email", "nick", "password", "options"}).
				AddRow("admin@cloudreve.org", "admin", "CKLmDKa1C9SD64vU:76adadd4fd4bad86959155f6f7bc8993c94e7adf", "{}"),
			policyRows: sqlmock.NewRows([]string{"name", "type", "options"}).
				AddRow("默认上传策略", "local", "{\"op_name\":\"123\"}"),
			reqBody:  `{"userName":"admin@cloudreve.org","captchaCode":"captchaCode","Password":"admin123"}`,
			expected: serializer.Err(401, "用户邮箱或密码错误", nil),
		},
		//邮箱格式不正确
		{
			reqBody:  `{"userName":"admin@cloudreve","captchaCode":"captchaCode","Password":"admin123"}`,
			expected: serializer.Err(40001, "邮箱格式不正确", errors.New("Key: 'UserLoginService.UserName' Error:Field validation for 'UserName' failed on the 'email' tag")),
		},
		// 用户被Ban
		{
			settingRows: sqlmock.NewRows([]string{"name", "value", "type"}).
				AddRow("login_captcha", "0", "login"),
			userRows: sqlmock.NewRows([]string{"email", "nick", "password", "options", "status"}).
				AddRow("admin@cloudreve.org", "admin", "CKLmDKa1C9SD64vU:76adadd4fd4bad86959155f6f7bc8993c94e7adf", "{}", model.Baned),
			policyRows: sqlmock.NewRows([]string{"name", "type", "options"}).
				AddRow("默认上传策略", "local", "{\"op_name\":\"123\"}"),
			reqBody:  `{"userName":"admin@cloudreve.org","captchaCode":"captchaCode","Password":"admin"}`,
			expected: serializer.Err(403, "该账号已被封禁", nil),
		},
		// 用户未激活
		{
			settingRows: sqlmock.NewRows([]string{"name", "value", "type"}).
				AddRow("login_captcha", "0", "login"),
			userRows: sqlmock.NewRows([]string{"email", "nick", "password", "options", "status"}).
				AddRow("admin@cloudreve.org", "admin", "CKLmDKa1C9SD64vU:76adadd4fd4bad86959155f6f7bc8993c94e7adf", "{}", model.NotActivicated),
			policyRows: sqlmock.NewRows([]string{"name", "type", "options"}).
				AddRow("默认上传策略", "local", "{\"op_name\":\"123\"}"),
			reqBody:  `{"userName":"admin@cloudreve.org","captchaCode":"captchaCode","Password":"admin"}`,
			expected: serializer.Err(403, "该账号未激活", nil),
		},
	}

	for k, testCase := range testCases {
		if testCase.settingRows != nil {
			mock.ExpectQuery("^SELECT (.+)").WillReturnRows(testCase.settingRows)
		}
		if testCase.userRows != nil {
			mock.ExpectQuery("^SELECT (.+)").WillReturnRows(testCase.userRows)
		}
		if testCase.policyRows != nil {
			mock.ExpectQuery("^SELECT \\* FROM `(.+)` WHERE `(.+)`\\.`deleted_at` IS NULL AND \\(\\(`policies`.`id` = 1\\)\\)(.+)$").WillReturnRows(testCase.policyRows)
		}
		req, _ := http.NewRequest(
			"POST",
			"/Api/V3/User/Session",
			bytes.NewReader([]byte(testCase.reqBody)),
		)
		router.ServeHTTP(w, req)

		asserts.Equal(200, w.Code)
		expectedJSON, _ := json.Marshal(testCase.expected)
		asserts.JSONEq(string(expectedJSON), w.Body.String(), "测试用例：%d", k)

		w.Body.Reset()
		asserts.NoError(mock.ExpectationsWereMet())
		model.ClearCache()
	}

}

func TestSessionAuthCheck(t *testing.T) {
	asserts := assert.New(t)
	router := InitRouter()
	w := httptest.NewRecorder()

	mock.ExpectQuery("^SELECT (.+)").WillReturnRows(sqlmock.NewRows([]string{"email", "nick", "password", "options"}).
		AddRow("admin@cloudreve.org", "admin", "CKLmDKa1C9SD64vU:76adadd4fd4bad86959155f6f7bc8993c94e7adf", "{}"))
	expectedUser, _ := model.GetUserByID(1)

	testCases := []struct {
		userRows    *sqlmock.Rows
		sessionMock map[string]interface{}
		contextMock map[string]interface{}
		expected    interface{}
	}{
		// 未登录
		{
			expected: serializer.CheckLogin(),
		},
		// 登录正常
		{
			userRows: sqlmock.NewRows([]string{"email", "nick", "password", "options"}).
				AddRow("admin@cloudreve.org", "admin", "CKLmDKa1C9SD64vU:76adadd4fd4bad86959155f6f7bc8993c94e7adf", "{}"),
			sessionMock: map[string]interface{}{"user_id": 1},
			expected:    serializer.BuildUserResponse(expectedUser),
		},
		// UID不存在
		{
			userRows:    sqlmock.NewRows([]string{"email", "nick", "password", "options"}),
			sessionMock: map[string]interface{}{"user_id": -1},
			expected:    serializer.CheckLogin(),
		},
	}

	for _, testCase := range testCases {
		req, _ := http.NewRequest(
			"GET",
			"/Api/V3/User/Me",
			nil,
		)
		if testCase.userRows != nil {
			mock.ExpectQuery("^SELECT (.+)").WillReturnRows(testCase.userRows)
		}
		middleware.ContextMock = testCase.contextMock
		middleware.SessionMock = testCase.sessionMock
		router.ServeHTTP(w, req)
		expectedJSON, _ := json.Marshal(testCase.expected)

		asserts.Equal(200, w.Code)
		asserts.JSONEq(string(expectedJSON), w.Body.String())
		asserts.NoError(mock.ExpectationsWereMet())

		w.Body.Reset()
	}

}