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) } }