feat: publish sanitized custom.97 source

2026-10-27 20:43:39 +08:00
1 parent 8e6ed0317b
commit 86c0276156
@@ -0,0 +1,113 @@
+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)
+		})
+	}
+}