easyai
3056cf8fca
ci / verify (pull_request) Successful in 15m34s
新增图片矢量化、视频超分、每日用量、计价与任务隔离能力，并通过环境变量解析平台凭据。 已通过 Go 全量门禁、迁移检查、镜像构建以及 Vectorizer 五格式和 Topaz 3 秒视频真实 DEV 验收。
80 lines
3.4 KiB
Go
80 lines
3.4 KiB
Go
package store
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"os"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
|
)
|
|
|
|
func TestListDailyTokenUsageScopesAPIKeyAndAppliesIANATimezone(t *testing.T) {
|
|
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
|
|
if databaseURL == "" {
|
|
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run daily usage PostgreSQL integration tests")
|
|
}
|
|
ctx := context.Background()
|
|
db, err := Connect(ctx, databaseURL)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer db.Close()
|
|
var databaseName string
|
|
if err := db.pool.QueryRow(ctx, `SELECT current_database()`).Scan(&databaseName); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if !strings.Contains(strings.ToLower(databaseName), "test") {
|
|
t.Fatalf("refusing to use non-test database %q", databaseName)
|
|
}
|
|
|
|
suffix := fmt.Sprint(time.Now().UnixNano())
|
|
var userID, otherUserID string
|
|
if err := db.pool.QueryRow(ctx, `INSERT INTO gateway_users (user_key, username) VALUES ($1, $2) RETURNING id::text`, "daily-user-"+suffix, "daily-user-"+suffix).Scan(&userID); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := db.pool.QueryRow(ctx, `INSERT INTO gateway_users (user_key, username) VALUES ($1, $2) RETURNING id::text`, "daily-other-"+suffix, "daily-other-"+suffix).Scan(&otherUserID); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
t.Cleanup(func() {
|
|
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE gateway_user_id IN ($1::uuid, $2::uuid)`, userID, otherUserID)
|
|
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_users WHERE id IN ($1::uuid, $2::uuid)`, userID, otherUserID)
|
|
})
|
|
insertTask := func(ownerID, keyID, createdAt string, totalTokens int, points float64) {
|
|
t.Helper()
|
|
_, err := db.pool.Exec(ctx, `
|
|
INSERT INTO gateway_tasks (kind, user_id, gateway_user_id, api_key_id, model, status, usage, final_charge_amount, created_at, finished_at)
|
|
VALUES ('chat.completions', $1::text, ($1::text)::uuid, $2::text, 'daily-test-model', 'succeeded', jsonb_build_object('totalTokens', $3::bigint, 'inputTokens', $3::bigint - 1, 'outputTokens', 1), $4::numeric, $5::timestamptz, $5::timestamptz)`, ownerID, keyID, totalTokens, points, createdAt)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
insertTask(userID, "daily-key-a-"+suffix, "2026-07-20T15:30:00Z", 10, 1)
|
|
insertTask(userID, "daily-key-a-"+suffix, "2026-07-21T16:30:00Z", 20, 2)
|
|
insertTask(userID, "daily-key-b-"+suffix, "2026-07-20T16:30:00Z", 100, 10)
|
|
insertTask(otherUserID, "daily-key-a-"+suffix, "2026-07-20T16:30:00Z", 999, 99)
|
|
|
|
location, err := time.LoadLocation("Asia/Shanghai")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
from := time.Date(2026, 7, 20, 0, 0, 0, 0, location)
|
|
to := time.Date(2026, 7, 22, 0, 0, 0, 0, location)
|
|
keyItems, err := db.ListDailyTokenUsage(ctx, &auth.User{GatewayUserID: userID, APIKeyID: "daily-key-a-" + suffix}, from, to, "Asia/Shanghai")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(keyItems) != 2 || keyItems[0].Date != "2026-07-20" || keyItems[0].TotalTokens != 10 || keyItems[1].Date != "2026-07-22" || keyItems[1].TotalTokens != 20 {
|
|
t.Fatalf("API key usage leaked another key/user or ignored timezone: %+v", keyItems)
|
|
}
|
|
userItems, err := db.ListDailyTokenUsage(ctx, &auth.User{GatewayUserID: userID}, from, to, "Asia/Shanghai")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(userItems) != 3 || userItems[1].Date != "2026-07-21" || userItems[1].TotalTokens != 100 {
|
|
t.Fatalf("JWT user usage did not include all owned API keys: %+v", userItems)
|
|
}
|
|
}
|