From 6ddd5bdf4e762e5c43f7c10ad9394c3287fdfadd Mon Sep 17 00:00:00 2001
From: ty <786220806@qq.com>
Date: Sun, 16 Mar 2025 15:34:50 +0800
Subject: [PATCH] chore: enhance test data generation with user identification

---
 scripts/generate-test-data.ts | 46 ++++++++++++++++++++++++++++++-----
 wrangler.cleanup.example.json |  2 +-
 wrangler.email.example.json   |  2 +-
 wrangler.example.json         |  2 +-
 4 files changed, 43 insertions(+), 9 deletions(-)

diff --git a/scripts/generate-test-data.ts b/scripts/generate-test-data.ts
index ac2216e..58f41f4 100644
--- a/scripts/generate-test-data.ts
+++ b/scripts/generate-test-data.ts
@@ -1,8 +1,7 @@
 import { drizzle } from 'drizzle-orm/d1'
-import { emails, messages } from '../app/lib/schema'
+import { emails, messages, users } from '../app/lib/schema'
 import { nanoid } from 'nanoid'
-
-const TEST_USER_ID = '4e4c1d5d-a3c9-407a-8808-2a2424b38c62'
+import { eq } from 'drizzle-orm'
 
 interface Env {
   DB: D1Database
@@ -12,16 +11,44 @@ const MAX_EMAIL_COUNT = 5
 const MAX_MESSAGE_COUNT = 50
 const BATCH_SIZE = 10 // SQLite 变量限制
 
-async function generateTestData(env: Env) {
+async function getUserId(db: ReturnType<typeof drizzle>, identifier: string): Promise<string | null> {
+  // 尝试通过 email 查找用户
+  let user = await db
+    .select()
+    .from(users)
+    .where(eq(users.email, identifier))
+    .limit(1)
+    .then(rows => rows[0])
+
+  // 如果没找到，尝试通过 username 查找
+  if (!user) {
+    user = await db
+      .select()
+      .from(users)
+      .where(eq(users.username, identifier))
+      .limit(1)
+      .then(rows => rows[0])
+  }
+
+  return user?.id || null
+}
+
+async function generateTestData(env: Env, userIdentifier: string) {
   const db = drizzle(env.DB)
   const now = new Date()
 
   try {
+    // 获取用户 ID
+    const userId = await getUserId(db, userIdentifier)
+    if (!userId) {
+      throw new Error(`未找到用户: ${userIdentifier}`)
+    }
+
     // 生成测试邮箱
     const testEmails = Array.from({ length: MAX_EMAIL_COUNT }).map(() => ({
       id: crypto.randomUUID(),
       address: `${nanoid(6)}@moemail.app`,
-      userId: TEST_USER_ID,
+      userId: userId,
       createdAt: now,
       expiresAt: new Date(now.getTime() + 24 * 60 * 60 * 1000),
     }))
@@ -66,7 +93,14 @@ async function generateTestData(env: Env) {
 export default {
   async fetch(request: Request, env: Env) {
     if (request.method === 'GET') {
-      await generateTestData(env)
+      const url = new URL(request.url)
+      const userIdentifier = url.searchParams.get('user')
+      
+      if (!userIdentifier) {
+        return new Response('Missing user parameter', { status: 400 })
+      }
+
+      await generateTestData(env, userIdentifier)
       return new Response('Test data generated successfully', { status: 200 })
     }
     return new Response('Method not allowed', { status: 405 })
diff --git a/wrangler.cleanup.example.json b/wrangler.cleanup.example.json
index 5ca3a68..439fe54 100644
--- a/wrangler.cleanup.example.json
+++ b/wrangler.cleanup.example.json
@@ -11,7 +11,7 @@
     {
       "binding": "DB",
       "migrations_dir": "drizzle",
-      "database_name": "${DATABASE_NAME}",
+      "database_name": "moemail",
       "database_id": "${DATABASE_ID}"
     }
   ]
diff --git a/wrangler.email.example.json b/wrangler.email.example.json
index dadd49e..03de89d 100644
--- a/wrangler.email.example.json
+++ b/wrangler.email.example.json
@@ -8,7 +8,7 @@
     {
       "binding": "DB",
       "migrations_dir": "drizzle",
-      "database_name": "${DATABASE_NAME}",
+      "database_name": "moemail",
       "database_id": "${DATABASE_ID}"
     }
   ]
diff --git a/wrangler.example.json b/wrangler.example.json
index 2c6bb08..a34d017 100644
--- a/wrangler.example.json
+++ b/wrangler.example.json
@@ -7,7 +7,7 @@
   "d1_databases": [
     {
       "binding": "DB",
-      "database_name": "${DATABASE_NAME}",
+      "database_name": "moemail",
       "database_id": "${DATABASE_ID}",
       "migrations_dir": "drizzle"
     }