feat: 测试数据生成脚本

scripts/generate-test-data.js

@@ -0,0 +1,284 @@
+#!/usr/bin/env node
+
+/**
+ * 历史数据生成脚本
+ * 用于测试不同时间范围的Token统计功能
+ * 
+ * 使用方法：
+ * node scripts/generate-test-data.js [--clean]
+ * 
+ * 选项：
+ * --clean: 清除所有测试数据
+ */
+
+const redis = require('../src/models/redis');
+const logger = require('../src/utils/logger');
+
+// 解析命令行参数
+const args = process.argv.slice(2);
+const shouldClean = args.includes('--clean');
+
+// 模拟的模型列表
+const models = [
+  'claude-sonnet-4-20250514',
+  'claude-3-5-sonnet-20241022',
+  'claude-3-5-haiku-20241022',
+  'claude-3-opus-20240229'
+];
+
+// 生成指定日期的数据
+async function generateDataForDate(apiKeyId, date, dayOffset) {
+  const client = redis.getClientSafe();
+  const dateStr = date.toISOString().split('T')[0];
+  const month = `${date.getFullYear()}-${String(date.getMonth() + 1).padStart(2, '0')}`;
+  
+  // 根据日期偏移量调整数据量（越近的日期数据越多）
+  const requestCount = Math.max(5, 20 - dayOffset * 2); // 5-20个请求
+  
+  logger.info(`📊 Generating ${requestCount} requests for ${dateStr}`);
+  
+  for (let i = 0; i < requestCount; i++) {
+    // 随机选择模型
+    const model = models[Math.floor(Math.random() * models.length)];
+    
+    // 生成随机Token数据
+    const inputTokens = Math.floor(Math.random() * 2000) + 500; // 500-2500
+    const outputTokens = Math.floor(Math.random() * 3000) + 1000; // 1000-4000
+    const cacheCreateTokens = Math.random() > 0.7 ? Math.floor(Math.random() * 1000) : 0; // 30%概率有缓存创建
+    const cacheReadTokens = Math.random() > 0.5 ? Math.floor(Math.random() * 500) : 0; // 50%概率有缓存读取
+    
+    const coreTokens = inputTokens + outputTokens;
+    const allTokens = inputTokens + outputTokens + cacheCreateTokens + cacheReadTokens;
+    
+    // 更新各种统计键
+    const totalKey = `usage:${apiKeyId}`;
+    const dailyKey = `usage:daily:${apiKeyId}:${dateStr}`;
+    const monthlyKey = `usage:monthly:${apiKeyId}:${month}`;
+    const modelDailyKey = `usage:model:daily:${model}:${dateStr}`;
+    const modelMonthlyKey = `usage:model:monthly:${model}:${month}`;
+    const keyModelDailyKey = `usage:${apiKeyId}:model:daily:${model}:${dateStr}`;
+    const keyModelMonthlyKey = `usage:${apiKeyId}:model:monthly:${model}:${month}`;
+    
+    await Promise.all([
+      // 总计数据
+      client.hincrby(totalKey, 'totalTokens', coreTokens),
+      client.hincrby(totalKey, 'totalInputTokens', inputTokens),
+      client.hincrby(totalKey, 'totalOutputTokens', outputTokens),
+      client.hincrby(totalKey, 'totalCacheCreateTokens', cacheCreateTokens),
+      client.hincrby(totalKey, 'totalCacheReadTokens', cacheReadTokens),
+      client.hincrby(totalKey, 'totalAllTokens', allTokens),
+      client.hincrby(totalKey, 'totalRequests', 1),
+      
+      // 每日统计
+      client.hincrby(dailyKey, 'tokens', coreTokens),
+      client.hincrby(dailyKey, 'inputTokens', inputTokens),
+      client.hincrby(dailyKey, 'outputTokens', outputTokens),
+      client.hincrby(dailyKey, 'cacheCreateTokens', cacheCreateTokens),
+      client.hincrby(dailyKey, 'cacheReadTokens', cacheReadTokens),
+      client.hincrby(dailyKey, 'allTokens', allTokens),
+      client.hincrby(dailyKey, 'requests', 1),
+      
+      // 每月统计
+      client.hincrby(monthlyKey, 'tokens', coreTokens),
+      client.hincrby(monthlyKey, 'inputTokens', inputTokens),
+      client.hincrby(monthlyKey, 'outputTokens', outputTokens),
+      client.hincrby(monthlyKey, 'cacheCreateTokens', cacheCreateTokens),
+      client.hincrby(monthlyKey, 'cacheReadTokens', cacheReadTokens),
+      client.hincrby(monthlyKey, 'allTokens', allTokens),
+      client.hincrby(monthlyKey, 'requests', 1),
+      
+      // 模型统计 - 每日
+      client.hincrby(modelDailyKey, 'totalInputTokens', inputTokens),
+      client.hincrby(modelDailyKey, 'totalOutputTokens', outputTokens),
+      client.hincrby(modelDailyKey, 'totalCacheCreateTokens', cacheCreateTokens),
+      client.hincrby(modelDailyKey, 'totalCacheReadTokens', cacheReadTokens),
+      client.hincrby(modelDailyKey, 'totalAllTokens', allTokens),
+      client.hincrby(modelDailyKey, 'requests', 1),
+      
+      // 模型统计 - 每月
+      client.hincrby(modelMonthlyKey, 'totalInputTokens', inputTokens),
+      client.hincrby(modelMonthlyKey, 'totalOutputTokens', outputTokens),
+      client.hincrby(modelMonthlyKey, 'totalCacheCreateTokens', cacheCreateTokens),
+      client.hincrby(modelMonthlyKey, 'totalCacheReadTokens', cacheReadTokens),
+      client.hincrby(modelMonthlyKey, 'totalAllTokens', allTokens),
+      client.hincrby(modelMonthlyKey, 'requests', 1),
+      
+      // API Key级别的模型统计 - 每日
+      // 同时存储带total前缀和不带前缀的字段，保持兼容性
+      client.hincrby(keyModelDailyKey, 'inputTokens', inputTokens),
+      client.hincrby(keyModelDailyKey, 'outputTokens', outputTokens),
+      client.hincrby(keyModelDailyKey, 'cacheCreateTokens', cacheCreateTokens),
+      client.hincrby(keyModelDailyKey, 'cacheReadTokens', cacheReadTokens),
+      client.hincrby(keyModelDailyKey, 'allTokens', allTokens),
+      client.hincrby(keyModelDailyKey, 'totalInputTokens', inputTokens),
+      client.hincrby(keyModelDailyKey, 'totalOutputTokens', outputTokens),
+      client.hincrby(keyModelDailyKey, 'totalCacheCreateTokens', cacheCreateTokens),
+      client.hincrby(keyModelDailyKey, 'totalCacheReadTokens', cacheReadTokens),
+      client.hincrby(keyModelDailyKey, 'totalAllTokens', allTokens),
+      client.hincrby(keyModelDailyKey, 'requests', 1),
+      
+      // API Key级别的模型统计 - 每月
+      client.hincrby(keyModelMonthlyKey, 'inputTokens', inputTokens),
+      client.hincrby(keyModelMonthlyKey, 'outputTokens', outputTokens),
+      client.hincrby(keyModelMonthlyKey, 'cacheCreateTokens', cacheCreateTokens),
+      client.hincrby(keyModelMonthlyKey, 'cacheReadTokens', cacheReadTokens),
+      client.hincrby(keyModelMonthlyKey, 'allTokens', allTokens),
+      client.hincrby(keyModelMonthlyKey, 'totalInputTokens', inputTokens),
+      client.hincrby(keyModelMonthlyKey, 'totalOutputTokens', outputTokens),
+      client.hincrby(keyModelMonthlyKey, 'totalCacheCreateTokens', cacheCreateTokens),
+      client.hincrby(keyModelMonthlyKey, 'totalCacheReadTokens', cacheReadTokens),
+      client.hincrby(keyModelMonthlyKey, 'totalAllTokens', allTokens),
+      client.hincrby(keyModelMonthlyKey, 'requests', 1),
+    ]);
+  }
+}
+
+// 清除测试数据
+async function cleanTestData() {
+  const client = redis.getClientSafe();
+  const apiKeyService = require('../src/services/apiKeyService');
+  
+  logger.info('🧹 Cleaning test data...');
+  
+  // 获取所有API Keys
+  const allKeys = await apiKeyService.getAllApiKeys();
+  
+  // 找出所有测试 API Keys
+  const testKeys = allKeys.filter(key => key.name && key.name.startsWith('Test API Key'));
+  
+  for (const testKey of testKeys) {
+    const apiKeyId = testKey.id;
+    
+    // 获取所有相关的键
+    const patterns = [
+      `usage:${apiKeyId}`,
+      `usage:daily:${apiKeyId}:*`,
+      `usage:monthly:${apiKeyId}:*`,
+      `usage:${apiKeyId}:model:daily:*`,
+      `usage:${apiKeyId}:model:monthly:*`
+    ];
+    
+    for (const pattern of patterns) {
+      const keys = await client.keys(pattern);
+      if (keys.length > 0) {
+        await client.del(...keys);
+        logger.info(`🗑️ Deleted ${keys.length} keys matching pattern: ${pattern}`);
+      }
+    }
+    
+    // 删除 API Key 本身
+    await apiKeyService.deleteApiKey(apiKeyId);
+    logger.info(`🗑️ Deleted test API Key: ${testKey.name} (${apiKeyId})`);
+  }
+  
+  // 清除模型统计
+  const modelPatterns = [
+    'usage:model:daily:*',
+    'usage:model:monthly:*'
+  ];
+  
+  for (const pattern of modelPatterns) {
+    const keys = await client.keys(pattern);
+    if (keys.length > 0) {
+      await client.del(...keys);
+      logger.info(`🗑️ Deleted ${keys.length} keys matching pattern: ${pattern}`);
+    }
+  }
+}
+
+// 主函数
+async function main() {
+  try {
+    await redis.connect();
+    logger.success('✅ Connected to Redis');
+    
+    // 创建测试API Keys
+    const apiKeyService = require('../src/services/apiKeyService');
+    let testApiKeys = [];
+    let createdKeys = [];
+    
+    // 总是创建新的测试 API Keys
+    logger.info('📝 Creating test API Keys...');
+    
+    for (let i = 1; i <= 3; i++) {
+      const newKey = await apiKeyService.generateApiKey({
+        name: `Test API Key ${i}`,
+        description: `Test key for historical data generation ${i}`,
+        tokenLimit: 10000000,
+        concurrencyLimit: 10,
+        rateLimitWindow: 60,
+        rateLimitRequests: 100
+      });
+      
+      testApiKeys.push(newKey.id);
+      createdKeys.push(newKey);
+      logger.success(`✅ Created test API Key: ${newKey.name} (${newKey.id})`);
+      logger.info(`   🔑 API Key: ${newKey.apiKey}`);
+    }
+    
+    if (shouldClean) {
+      await cleanTestData();
+      logger.success('✅ Test data cleaned successfully');
+      return;
+    }
+    
+    // 生成历史数据
+    const now = new Date();
+    
+    for (const apiKeyId of testApiKeys) {
+      logger.info(`\n🔄 Generating data for API Key: ${apiKeyId}`);
+      
+      // 生成过去30天的数据
+      for (let dayOffset = 0; dayOffset < 30; dayOffset++) {
+        const date = new Date(now);
+        date.setDate(date.getDate() - dayOffset);
+        
+        await generateDataForDate(apiKeyId, date, dayOffset);
+      }
+      
+      logger.success(`✅ Generated 30 days of historical data for API Key: ${apiKeyId}`);
+    }
+    
+    // 显示统计摘要
+    logger.info('\n📊 Test Data Summary:');
+    logger.info('='.repeat(60));
+    
+    for (const apiKeyId of testApiKeys) {
+      const totalKey = `usage:${apiKeyId}`;
+      const totalData = await redis.getClientSafe().hgetall(totalKey);
+      
+      if (totalData && Object.keys(totalData).length > 0) {
+        logger.info(`\nAPI Key: ${apiKeyId}`);
+        logger.info(`  Total Requests: ${totalData.totalRequests || 0}`);
+        logger.info(`  Total Tokens (Core): ${totalData.totalTokens || 0}`);
+        logger.info(`  Total Tokens (All): ${totalData.totalAllTokens || 0}`);
+        logger.info(`  Input Tokens: ${totalData.totalInputTokens || 0}`);
+        logger.info(`  Output Tokens: ${totalData.totalOutputTokens || 0}`);
+        logger.info(`  Cache Create Tokens: ${totalData.totalCacheCreateTokens || 0}`);
+        logger.info(`  Cache Read Tokens: ${totalData.totalCacheReadTokens || 0}`);
+      }
+    }
+    
+    logger.info('\n' + '='.repeat(60));
+    logger.success('\n✅ Test data generation completed!');
+    logger.info('\n📋 Created API Keys:');
+    for (const key of createdKeys) {
+      logger.info(`- ${key.name}: ${key.apiKey}`);
+    }
+    logger.info('\n💡 Tips:');
+    logger.info('- Check the admin panel to see the different time ranges');
+    logger.info('- Use --clean flag to remove all test data and API Keys');
+    logger.info('- The script generates more recent data to simulate real usage patterns');
+    
+  } catch (error) {
+    logger.error('❌ Error:', error);
+  } finally {
+    await redis.disconnect();
+  }
+}
+
+// 运行脚本
+main().catch(error => {
+  logger.error('💥 Unexpected error:', error);
+  process.exit(1);
+});