Files
gpt-plus-gpt/internal/handler/handler_test.go
2026-03-15 20:48:19 +08:00

519 lines
13 KiB
Go

package handler
import (
"bytes"
"encoding/json"
"fmt"
"net/http"
"time"
"net/http/httptest"
"os"
"testing"
"gpt-plus/internal/db"
"github.com/glebarez/sqlite"
"github.com/gin-gonic/gin"
"gorm.io/gorm"
"gorm.io/gorm/logger"
)
func init() {
gin.SetMode(gin.TestMode)
}
func setupTestEnv(t *testing.T) *gorm.DB {
t.Helper()
d, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{
Logger: logger.Default.LogMode(logger.Silent),
})
if err != nil {
t.Fatalf("open test db: %v", err)
}
if err := d.AutoMigrate(
&db.SystemConfig{}, &db.EmailRecord{}, &db.Task{}, &db.TaskLog{},
&db.CardCode{}, &db.Card{}, &db.Account{},
); err != nil {
t.Fatalf("migrate: %v", err)
}
db.DB = d
SetJWTSecret("test-secret-key")
os.Setenv("ADMIN_PASSWORD", "testpass")
if err := InitAdminPassword(); err != nil {
t.Fatalf("init password: %v", err)
}
// Reset rate limiter between tests
loginLimiter.mu.Lock()
loginLimiter.attempts = make(map[string][]time.Time)
loginLimiter.mu.Unlock()
return d
}
func jsonReq(method, path string, body interface{}) *http.Request {
var buf bytes.Buffer
if body != nil {
json.NewEncoder(&buf).Encode(body)
}
req := httptest.NewRequest(method, path, &buf)
req.Header.Set("Content-Type", "application/json")
return req
}
func doRequest(r *gin.Engine, req *http.Request) *httptest.ResponseRecorder {
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
return w
}
func loginAndGetCookie(t *testing.T, r *gin.Engine) *http.Cookie {
t.Helper()
req := jsonReq("POST", "/api/login", map[string]string{"password": "testpass"})
w := doRequest(r, req)
if w.Code != http.StatusOK {
t.Fatalf("login failed: %d %s", w.Code, w.Body.String())
}
for _, c := range w.Result().Cookies() {
if c.Name == "token" {
return c
}
}
t.Fatal("no token cookie in login response")
return nil
}
func authedReq(method, path string, body interface{}, cookie *http.Cookie) *http.Request {
req := jsonReq(method, path, body)
req.AddCookie(cookie)
return req
}
// --- Auth tests ---
func TestLoginSuccess(t *testing.T) {
setupTestEnv(t)
r := SetupRouter(false)
req := jsonReq("POST", "/api/login", map[string]string{"password": "testpass"})
w := doRequest(r, req)
if w.Code != http.StatusOK {
t.Fatalf("status = %d, want 200", w.Code)
}
var resp map[string]string
json.Unmarshal(w.Body.Bytes(), &resp)
if resp["message"] != "登录成功" {
t.Fatalf("message = %q", resp["message"])
}
found := false
for _, c := range w.Result().Cookies() {
if c.Name == "token" && c.HttpOnly {
found = true
}
}
if !found {
t.Fatal("no httpOnly token cookie set")
}
}
func TestLoginWrongPassword(t *testing.T) {
setupTestEnv(t)
r := SetupRouter(false)
req := jsonReq("POST", "/api/login", map[string]string{"password": "wrong"})
w := doRequest(r, req)
if w.Code != http.StatusUnauthorized {
t.Fatalf("status = %d, want 401", w.Code)
}
}
func TestLoginEmptyPassword(t *testing.T) {
setupTestEnv(t)
r := SetupRouter(false)
req := jsonReq("POST", "/api/login", map[string]string{})
w := doRequest(r, req)
if w.Code != http.StatusBadRequest {
t.Fatalf("status = %d, want 400", w.Code)
}
}
func TestLogout(t *testing.T) {
setupTestEnv(t)
r := SetupRouter(false)
req := jsonReq("POST", "/api/logout", nil)
w := doRequest(r, req)
if w.Code != http.StatusOK {
t.Fatalf("status = %d, want 200", w.Code)
}
for _, c := range w.Result().Cookies() {
if c.Name == "token" && c.MaxAge < 0 {
return
}
}
t.Fatal("token cookie not cleared")
}
func TestAuthMiddlewareRejectsUnauthenticated(t *testing.T) {
setupTestEnv(t)
r := SetupRouter(false)
req := jsonReq("GET", "/api/config", nil)
w := doRequest(r, req)
if w.Code != http.StatusUnauthorized {
t.Fatalf("status = %d, want 401", w.Code)
}
}
func TestAuthMiddlewareAcceptsValidToken(t *testing.T) {
setupTestEnv(t)
SeedDefaultConfigs()
r := SetupRouter(false)
cookie := loginAndGetCookie(t, r)
req := authedReq("GET", "/api/auth/check", nil, cookie)
w := doRequest(r, req)
if w.Code != http.StatusOK {
t.Fatalf("status = %d, want 200", w.Code)
}
}
// --- Config tests ---
func TestGetConfig(t *testing.T) {
setupTestEnv(t)
SeedDefaultConfigs()
r := SetupRouter(false)
cookie := loginAndGetCookie(t, r)
req := authedReq("GET", "/api/config", nil, cookie)
w := doRequest(r, req)
if w.Code != http.StatusOK {
t.Fatalf("status = %d, want 200, body: %s", w.Code, w.Body.String())
}
var resp struct {
Groups map[string][]db.SystemConfig `json:"groups"`
}
json.Unmarshal(w.Body.Bytes(), &resp)
if len(resp.Groups) == 0 {
t.Fatal("expected non-empty config groups")
}
if _, ok := resp.Groups["proxy"]; !ok {
t.Fatal("expected proxy group")
}
}
func TestUpdateConfig(t *testing.T) {
setupTestEnv(t)
SeedDefaultConfigs()
r := SetupRouter(false)
cookie := loginAndGetCookie(t, r)
body := map[string]interface{}{
"items": []map[string]string{
{"key": "proxy.url", "value": "socks5://1.2.3.4:1080"},
},
}
req := authedReq("PUT", "/api/config", body, cookie)
w := doRequest(r, req)
if w.Code != http.StatusOK {
t.Fatalf("status = %d, want 200, body: %s", w.Code, w.Body.String())
}
var cfg db.SystemConfig
db.GetDB().Where("key = ?", "proxy.url").First(&cfg)
if cfg.Value != "socks5://1.2.3.4:1080" {
t.Fatalf("config value = %q, want socks5://1.2.3.4:1080", cfg.Value)
}
}
// --- Card tests ---
func TestCardCRUD(t *testing.T) {
setupTestEnv(t)
SeedDefaultConfigs()
r := SetupRouter(false)
cookie := loginAndGetCookie(t, r)
// Add card
body := map[string]string{
"number": "4242424242424242", "cvc": "123",
"exp_month": "12", "exp_year": "2030", "country": "US",
}
req := authedReq("POST", "/api/cards", body, cookie)
w := doRequest(r, req)
if w.Code != http.StatusCreated {
t.Fatalf("add card: status = %d, body: %s", w.Code, w.Body.String())
}
// List cards
req = authedReq("GET", "/api/cards", nil, cookie)
w = doRequest(r, req)
if w.Code != http.StatusOK {
t.Fatalf("list cards: status = %d", w.Code)
}
var listResp struct {
Items []db.Card `json:"items"`
Total int64 `json:"total"`
}
json.Unmarshal(w.Body.Bytes(), &listResp)
if listResp.Total != 1 {
t.Fatalf("total = %d, want 1", listResp.Total)
}
if listResp.Items[0].NumberLast4 != "4242" {
t.Fatalf("last4 = %q, want 4242", listResp.Items[0].NumberLast4)
}
// Stats
req = authedReq("GET", "/api/cards/stats", nil, cookie)
w = doRequest(r, req)
if w.Code != http.StatusOK {
t.Fatalf("card stats: status = %d", w.Code)
}
// Delete
cardID := listResp.Items[0].ID
req = authedReq("DELETE", "/api/cards/"+itoa(cardID), nil, cookie)
w = doRequest(r, req)
if w.Code != http.StatusOK {
t.Fatalf("delete card: status = %d", w.Code)
}
}
func TestCardDuplicate(t *testing.T) {
setupTestEnv(t)
SeedDefaultConfigs()
r := SetupRouter(false)
cookie := loginAndGetCookie(t, r)
body := map[string]string{
"number": "5555555555554444", "cvc": "321",
"exp_month": "06", "exp_year": "2029", "country": "JP",
}
req := authedReq("POST", "/api/cards", body, cookie)
doRequest(r, req) // first
req = authedReq("POST", "/api/cards", body, cookie)
w := doRequest(r, req) // duplicate
if w.Code != http.StatusConflict {
t.Fatalf("duplicate card: status = %d, want 409", w.Code)
}
}
// --- Card Code tests ---
func TestCardCodeImport(t *testing.T) {
setupTestEnv(t)
SeedDefaultConfigs()
r := SetupRouter(false)
cookie := loginAndGetCookie(t, r)
body := map[string]interface{}{
"codes": []string{"CODE-001", "CODE-002", "CODE-003"},
}
req := authedReq("POST", "/api/card-codes/import", body, cookie)
w := doRequest(r, req)
if w.Code != http.StatusOK {
t.Fatalf("import: status = %d, body: %s", w.Code, w.Body.String())
}
var resp map[string]int
json.Unmarshal(w.Body.Bytes(), &resp)
if resp["imported"] != 3 {
t.Fatalf("imported = %d, want 3", resp["imported"])
}
// Stats
req = authedReq("GET", "/api/card-codes/stats", nil, cookie)
w = doRequest(r, req)
var stats db.CardCodeStats
json.Unmarshal(w.Body.Bytes(), &stats)
if stats.Unused != 3 {
t.Fatalf("unused = %d, want 3", stats.Unused)
}
}
// --- Dashboard test ---
func TestDashboard(t *testing.T) {
setupTestEnv(t)
SeedDefaultConfigs()
r := SetupRouter(false)
cookie := loginAndGetCookie(t, r)
req := authedReq("GET", "/api/dashboard", nil, cookie)
w := doRequest(r, req)
if w.Code != http.StatusOK {
t.Fatalf("dashboard: status = %d, body: %s", w.Code, w.Body.String())
}
}
// --- Email record tests ---
func TestEmailRecordStats(t *testing.T) {
d := setupTestEnv(t)
SeedDefaultConfigs()
r := SetupRouter(false)
cookie := loginAndGetCookie(t, r)
d.Create(&db.EmailRecord{Email: "a@test.com", Status: "used"})
d.Create(&db.EmailRecord{Email: "b@test.com", Status: "used_failed"})
req := authedReq("GET", "/api/email-records/stats", nil, cookie)
w := doRequest(r, req)
if w.Code != http.StatusOK {
t.Fatalf("status = %d", w.Code)
}
var stats db.EmailRecordStats
json.Unmarshal(w.Body.Bytes(), &stats)
if stats.Total != 2 || stats.Used != 1 || stats.UsedFailed != 1 {
t.Fatalf("stats = %+v", stats)
}
}
// --- Account tests ---
func TestAccountListAndDetail(t *testing.T) {
d := setupTestEnv(t)
SeedDefaultConfigs()
r := SetupRouter(false)
cookie := loginAndGetCookie(t, r)
d.Create(&db.Account{Email: "plus@test.com", Plan: "plus", Status: "active", AccountID: "acc-1"})
owner := db.Account{Email: "team@test.com", Plan: "team_owner", Status: "active", AccountID: "acc-2"}
d.Create(&owner)
d.Create(&db.Account{Email: "m1@test.com", Plan: "team_member", Status: "active", ParentID: &owner.ID})
// List (should only show plus + team_owner)
req := authedReq("GET", "/api/accounts", nil, cookie)
w := doRequest(r, req)
if w.Code != http.StatusOK {
t.Fatalf("status = %d", w.Code)
}
var listResp struct {
Items []db.Account `json:"items"`
Total int64 `json:"total"`
}
json.Unmarshal(w.Body.Bytes(), &listResp)
if listResp.Total != 2 {
t.Fatalf("total = %d, want 2 (team_member hidden)", listResp.Total)
}
// Detail with sub-accounts
req = authedReq("GET", "/api/accounts/"+itoa(owner.ID), nil, cookie)
w = doRequest(r, req)
if w.Code != http.StatusOK {
t.Fatalf("detail status = %d", w.Code)
}
var detail db.Account
json.Unmarshal(w.Body.Bytes(), &detail)
if len(detail.SubAccounts) != 1 {
t.Fatalf("sub_accounts = %d, want 1", len(detail.SubAccounts))
}
}
func TestAccountUpdateNote(t *testing.T) {
d := setupTestEnv(t)
SeedDefaultConfigs()
r := SetupRouter(false)
cookie := loginAndGetCookie(t, r)
d.Create(&db.Account{Email: "note@test.com", Plan: "plus", Status: "active"})
req := authedReq("PUT", "/api/accounts/1/note", map[string]string{"note": "VIP"}, cookie)
w := doRequest(r, req)
if w.Code != http.StatusOK {
t.Fatalf("status = %d", w.Code)
}
var acct db.Account
d.First(&acct, 1)
if acct.Note != "VIP" {
t.Fatalf("note = %q, want VIP", acct.Note)
}
}
func TestAccountSearch(t *testing.T) {
d := setupTestEnv(t)
SeedDefaultConfigs()
r := SetupRouter(false)
cookie := loginAndGetCookie(t, r)
d.Create(&db.Account{Email: "alice@test.com", Plan: "plus", Status: "active"})
d.Create(&db.Account{Email: "bob@test.com", Plan: "plus", Status: "active"})
req := authedReq("GET", "/api/accounts?search=alice", nil, cookie)
w := doRequest(r, req)
var resp struct {
Total int64 `json:"total"`
}
json.Unmarshal(w.Body.Bytes(), &resp)
if resp.Total != 1 {
t.Fatalf("search total = %d, want 1", resp.Total)
}
}
// --- Task tests ---
func TestTaskCreateAndList(t *testing.T) {
setupTestEnv(t)
SeedDefaultConfigs()
r := SetupRouter(false)
cookie := loginAndGetCookie(t, r)
body := map[string]interface{}{"type": "plus", "count": 5}
req := authedReq("POST", "/api/tasks", body, cookie)
w := doRequest(r, req)
if w.Code != http.StatusCreated {
t.Fatalf("create task: status = %d, body: %s", w.Code, w.Body.String())
}
var created db.Task
json.Unmarshal(w.Body.Bytes(), &created)
if created.ID == "" || created.Status != "pending" || created.TotalCount != 5 {
t.Fatalf("created = %+v", created)
}
// List
req = authedReq("GET", "/api/tasks", nil, cookie)
w = doRequest(r, req)
if w.Code != http.StatusOK {
t.Fatalf("list: status = %d", w.Code)
}
// Delete
req = authedReq("DELETE", "/api/tasks/"+created.ID, nil, cookie)
w = doRequest(r, req)
if w.Code != http.StatusOK {
t.Fatalf("delete: status = %d, body: %s", w.Code, w.Body.String())
}
}
func TestTaskCreateInvalidType(t *testing.T) {
setupTestEnv(t)
SeedDefaultConfigs()
r := SetupRouter(false)
cookie := loginAndGetCookie(t, r)
body := map[string]interface{}{"type": "invalid", "count": 1}
req := authedReq("POST", "/api/tasks", body, cookie)
w := doRequest(r, req)
if w.Code != http.StatusBadRequest {
t.Fatalf("status = %d, want 400", w.Code)
}
}
func itoa(n uint) string {
return fmt.Sprintf("%d", n)
}