114 lines
4.7 KiB
Go
114 lines
4.7 KiB
Go
package controller
|
|
|
|
import (
|
|
"fmt"
|
|
"net/http/httptest"
|
|
"os"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/model"
|
|
"github.com/QuantumNous/new-api/setting/operation_setting"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/glebarez/sqlite"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/tidwall/gjson"
|
|
"gorm.io/driver/mysql"
|
|
"gorm.io/driver/postgres"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/schema"
|
|
)
|
|
|
|
func TestBotBalanceQuery(t *testing.T) {
|
|
key := strings.Repeat("b", 48)
|
|
t.Setenv("BOT_BALANCE_KEY", key)
|
|
oldUnit, oldRate := common.QuotaPerUnit, operation_setting.USDExchangeRate
|
|
common.QuotaPerUnit, operation_setting.USDExchangeRate = 500000, 7.3
|
|
t.Cleanup(func() { common.QuotaPerUnit, operation_setting.USDExchangeRate = oldUnit, oldRate })
|
|
for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
|
|
t.Run(dialect, func(t *testing.T) {
|
|
var driver gorm.Dialector = sqlite.Open(":memory:")
|
|
if dialect == "mysql" {
|
|
require.NotEmpty(t, os.Getenv("TEST_MYSQL_DSN"))
|
|
driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
|
|
}
|
|
if dialect == "postgres" {
|
|
require.NotEmpty(t, os.Getenv("TEST_POSTGRES_DSN"))
|
|
driver = postgres.Open(os.Getenv("TEST_POSTGRES_DSN"))
|
|
}
|
|
db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("bot_%d_", time.Now().UnixNano())}})
|
|
require.NoError(t, err)
|
|
conn, err := db.DB()
|
|
require.NoError(t, err)
|
|
conn.SetMaxOpenConns(1)
|
|
oldDB := model.DB
|
|
model.DB = db
|
|
t.Cleanup(func() { require.NoError(t, db.Migrator().DropTable(&model.User{})); model.DB = oldDB; _ = conn.Close() })
|
|
require.NoError(t, db.AutoMigrate(&model.User{}))
|
|
rows := []model.User{
|
|
{Id: 1, Username: "bot1", Email: "a@example.com", Status: 1, Quota: 250000},
|
|
{Id: 2, Username: "bot2", Email: "dup@example.com", Status: 1, Quota: 0},
|
|
{Id: 3, Username: "bot3", Email: "dup@example.com", Status: 1, Quota: -500000},
|
|
{Id: 4, Username: "bot4", Email: "disabled@example.com", Status: 2, Quota: 900000},
|
|
{Id: 5, Username: "bot5", Email: "gone@example.com", Status: 1},
|
|
}
|
|
for i := range rows {
|
|
rows[i].AffCode = fmt.Sprintf("bot%d", i)
|
|
}
|
|
require.NoError(t, db.Create(&rows).Error)
|
|
require.NoError(t, db.Delete(&rows[4]).Error)
|
|
r := gin.New()
|
|
r.POST("/api/bot/balances", BotBalanceQuery)
|
|
for _, tc := range []struct {
|
|
name, body, token string
|
|
status int
|
|
code string
|
|
}{
|
|
{"missing key", `{}`, "", 401, "unauthorized"},
|
|
{"wrong key", `{}`, "Bearer wrong", 401, "unauthorized"},
|
|
{"empty", `{}`, "Bearer " + key, 400, "invalid_request"},
|
|
{"both", `{"email":"a@example.com","ids":[1]}`, "Bearer " + key, 400, "invalid_request"},
|
|
{"bad id", `{"ids":[0]}`, "Bearer " + key, 400, "invalid_request"},
|
|
{"bad email", `{"email":"invalid"}`, "Bearer " + key, 400, "invalid_email"},
|
|
{"deleted", `{"email":"gone@example.com"}`, "Bearer " + key, 409, "not_found"},
|
|
{"duplicate", `{"email":"dup@example.com"}`, "Bearer " + key, 409, "ambiguous_email"},
|
|
{"exact not substring", `{"email":"aa@example.com"}`, "Bearer " + key, 409, "not_found"},
|
|
{"normalized", `{"email":" A@EXAMPLE.COM "}`, "Bearer " + key, 200, ""},
|
|
{"batch", `{"ids":[1,2,3,4,5,999]}`, "Bearer " + key, 200, ""},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest("POST", "/api/bot/balances", strings.NewReader(tc.body))
|
|
req.Header.Set("Authorization", tc.token)
|
|
req.Header.Set("Content-Type", "application/json")
|
|
r.ServeHTTP(w, req)
|
|
require.Equal(t, tc.status, w.Code, w.Body.String())
|
|
assert.Equal(t, "no-store", w.Header().Get("Cache-Control"))
|
|
assert.Equal(t, tc.code, gjson.Get(w.Body.String(), "code").String())
|
|
if tc.status == 200 {
|
|
assert.NotContains(t, w.Body.String(), "email")
|
|
assert.NotContains(t, w.Body.String(), "quota")
|
|
assert.NotContains(t, w.Body.String(), "password")
|
|
assert.Equal(t, "3.65", gjson.Get(w.Body.String(), "data.0.balance_cny").String())
|
|
if tc.name == "batch" {
|
|
assert.Equal(t, int64(4), gjson.Get(w.Body.String(), "data.#").Int())
|
|
assert.Equal(t, "0", gjson.Get(w.Body.String(), "data.1.balance_cny").String())
|
|
assert.Equal(t, "-7.3", gjson.Get(w.Body.String(), "data.2.balance_cny").String())
|
|
assert.False(t, gjson.Get(w.Body.String(), "data.3.active").Bool())
|
|
assert.False(t, gjson.Get(w.Body.String(), "data.3.balance_cny").Exists())
|
|
}
|
|
}
|
|
})
|
|
}
|
|
t.Setenv("BOT_BALANCE_KEY", "")
|
|
w := httptest.NewRecorder()
|
|
r.ServeHTTP(w, httptest.NewRequest("POST", "/api/bot/balances", strings.NewReader(`{}`)))
|
|
assert.Equal(t, 401, w.Code)
|
|
t.Setenv("BOT_BALANCE_KEY", key)
|
|
})
|
|
}
|
|
}
|