/** * 分组调度功能测试脚本 * 用于测试账户分组管理和调度逻辑的正确性 */ require('dotenv').config() const { v4: uuidv4 } = require('uuid') const redis = require('../src/models/redis') const accountGroupService = require('../src/services/accountGroupService') const claudeAccountService = require('../src/services/claudeAccountService') const claudeConsoleAccountService = require('../src/services/claudeConsoleAccountService') const apiKeyService = require('../src/services/apiKeyService') const unifiedClaudeScheduler = require('../src/services/unifiedClaudeScheduler') // 测试配置 const TEST_PREFIX = 'test_group_' const CLEANUP_ON_FINISH = true // 测试完成后是否清理数据 // 测试数据存储 const testData = { groups: [], accounts: [], apiKeys: [] } // 颜色输出 const colors = { green: '\x1b[32m', red: '\x1b[31m', yellow: '\x1b[33m', blue: '\x1b[34m', reset: '\x1b[0m' } function log(message, type = 'info') { const color = { success: colors.green, error: colors.red, warning: colors.yellow, info: colors.blue }[type] || colors.reset console.log(`${color}${message}${colors.reset}`) } async function sleep(ms) { return new Promise((resolve) => setTimeout(resolve, ms)) } // 清理测试数据 async function cleanup() { log('\n🧹 清理测试数据...', 'info') // 删除测试API Keys for (const apiKey of testData.apiKeys) { try { await apiKeyService.deleteApiKey(apiKey.id) log(`✅ 删除测试API Key: ${apiKey.name}`, 'success') } catch (error) { log(`❌ 删除API Key失败: ${error.message}`, 'error') } } // 删除测试账户 for (const account of testData.accounts) { try { if (account.type === 'claude') { await claudeAccountService.deleteAccount(account.id) } else if (account.type === 'claude-console') { await claudeConsoleAccountService.deleteAccount(account.id) } log(`✅ 删除测试账户: ${account.name}`, 'success') } catch (error) { log(`❌ 删除账户失败: ${error.message}`, 'error') } } // 删除测试分组 for (const group of testData.groups) { try { await accountGroupService.deleteGroup(group.id) log(`✅ 删除测试分组: ${group.name}`, 'success') } catch (error) { // 可能因为还有成员而删除失败,先移除所有成员 if (error.message.includes('分组内还有账户')) { const members = await accountGroupService.getGroupMembers(group.id) for (const memberId of members) { await accountGroupService.removeAccountFromGroup(memberId, group.id) } // 重试删除 await accountGroupService.deleteGroup(group.id) log(`✅ 删除测试分组: ${group.name} (清空成员后)`, 'success') } else { log(`❌ 删除分组失败: ${error.message}`, 'error') } } } } // 测试1: 创建分组 async function test1_createGroups() { log('\n📝 测试1: 创建账户分组', 'info') try { // 创建Claude分组 const claudeGroup = await accountGroupService.createGroup({ name: `${TEST_PREFIX}Claude组`, platform: 'claude', description: '测试用Claude账户分组' }) testData.groups.push(claudeGroup) log(`✅ 创建Claude分组成功: ${claudeGroup.name} (ID: ${claudeGroup.id})`, 'success') // 创建Gemini分组 const geminiGroup = await accountGroupService.createGroup({ name: `${TEST_PREFIX}Gemini组`, platform: 'gemini', description: '测试用Gemini账户分组' }) testData.groups.push(geminiGroup) log(`✅ 创建Gemini分组成功: ${geminiGroup.name} (ID: ${geminiGroup.id})`, 'success') // 验证分组信息 const allGroups = await accountGroupService.getAllGroups() const testGroups = allGroups.filter((g) => g.name.startsWith(TEST_PREFIX)) if (testGroups.length === 2) { log(`✅ 分组创建验证通过,共创建 ${testGroups.length} 个测试分组`, 'success') } else { throw new Error(`分组数量不正确,期望2个,实际${testGroups.length}个`) } } catch (error) { log(`❌ 测试1失败: ${error.message}`, 'error') throw error } } // 测试2: 创建账户并添加到分组 async function test2_createAccountsAndAddToGroup() { log('\n📝 测试2: 创建账户并添加到分组', 'info') try { const claudeGroup = testData.groups.find((g) => g.platform === 'claude') // 创建Claude OAuth账户 const claudeAccount1 = await claudeAccountService.createAccount({ name: `${TEST_PREFIX}Claude账户1`, email: 'test1@example.com', refreshToken: 'test_refresh_token_1', accountType: 'group' }) testData.accounts.push({ ...claudeAccount1, type: 'claude' }) log(`✅ 创建Claude OAuth账户1成功: ${claudeAccount1.name}`, 'success') const claudeAccount2 = await claudeAccountService.createAccount({ name: `${TEST_PREFIX}Claude账户2`, email: 'test2@example.com', refreshToken: 'test_refresh_token_2', accountType: 'group' }) testData.accounts.push({ ...claudeAccount2, type: 'claude' }) log(`✅ 创建Claude OAuth账户2成功: ${claudeAccount2.name}`, 'success') // 创建Claude Console账户 const consoleAccount = await claudeConsoleAccountService.createAccount({ name: `${TEST_PREFIX}Console账户`, apiUrl: 'https://api.example.com', apiKey: 'test_api_key', accountType: 'group' }) testData.accounts.push({ ...consoleAccount, type: 'claude-console' }) log(`✅ 创建Claude Console账户成功: ${consoleAccount.name}`, 'success') // 添加账户到分组 await accountGroupService.addAccountToGroup(claudeAccount1.id, claudeGroup.id, 'claude') log('✅ 添加账户1到分组成功', 'success') await accountGroupService.addAccountToGroup(claudeAccount2.id, claudeGroup.id, 'claude') log('✅ 添加账户2到分组成功', 'success') await accountGroupService.addAccountToGroup(consoleAccount.id, claudeGroup.id, 'claude') log('✅ 添加Console账户到分组成功', 'success') // 验证分组成员 const members = await accountGroupService.getGroupMembers(claudeGroup.id) if (members.length === 3) { log(`✅ 分组成员验证通过,共有 ${members.length} 个成员`, 'success') } else { throw new Error(`分组成员数量不正确,期望3个,实际${members.length}个`) } } catch (error) { log(`❌ 测试2失败: ${error.message}`, 'error') throw error } } // 测试3: 平台一致性验证 async function test3_platformConsistency() { log('\n📝 测试3: 平台一致性验证', 'info') try { const geminiGroup = testData.groups.find((g) => g.platform === 'gemini') // 尝试将Claude账户添加到Gemini分组(应该失败) const claudeAccount = testData.accounts.find((a) => a.type === 'claude') try { await accountGroupService.addAccountToGroup(claudeAccount.id, geminiGroup.id, 'claude') throw new Error('平台验证失败:Claude账户不应该能添加到Gemini分组') } catch (error) { if (error.message.includes('平台与分组平台不匹配')) { log(`✅ 平台一致性验证通过:${error.message}`, 'success') } else { throw error } } } catch (error) { log(`❌ 测试3失败: ${error.message}`, 'error') throw error } } // 测试4: API Key绑定分组 async function test4_apiKeyBindGroup() { log('\n📝 测试4: API Key绑定分组', 'info') try { const claudeGroup = testData.groups.find((g) => g.platform === 'claude') // 创建绑定到分组的API Key const apiKey = await apiKeyService.generateApiKey({ name: `${TEST_PREFIX}API Key`, description: '测试分组调度的API Key', claudeAccountId: `group:${claudeGroup.id}`, permissions: 'claude' }) testData.apiKeys.push(apiKey) log(`✅ 创建API Key成功: ${apiKey.name} (绑定到分组: ${claudeGroup.name})`, 'success') // 验证API Key信息 const keyInfo = await redis.getApiKey(apiKey.id) if (keyInfo && keyInfo.claudeAccountId === `group:${claudeGroup.id}`) { log('✅ API Key分组绑定验证通过', 'success') } else { throw new Error('API Key分组绑定信息不正确') } } catch (error) { log(`❌ 测试4失败: ${error.message}`, 'error') throw error } } // 测试5: 分组调度负载均衡 async function test5_groupSchedulingLoadBalance() { log('\n📝 测试5: 分组调度负载均衡', 'info') try { const apiKey = testData.apiKeys[0] // 记录每个账户被选中的次数 const selectionCount = {} const totalSelections = 30 for (let i = 0; i < totalSelections; i++) { // 模拟不同的会话 const sessionHash = uuidv4() const result = await unifiedClaudeScheduler.selectAccountForApiKey( { id: apiKey.id, claudeAccountId: apiKey.claudeAccountId, name: apiKey.name }, sessionHash ) if (!selectionCount[result.accountId]) { selectionCount[result.accountId] = 0 } selectionCount[result.accountId]++ // 短暂延迟,模拟真实请求间隔 await sleep(50) } // 分析选择分布 log(`\n📊 负载均衡分布统计 (共${totalSelections}次选择):`, 'info') const accounts = Object.keys(selectionCount) for (const accountId of accounts) { const count = selectionCount[accountId] const percentage = ((count / totalSelections) * 100).toFixed(1) const accountInfo = testData.accounts.find((a) => a.id === accountId) log(` ${accountInfo.name}: ${count}次 (${percentage}%)`, 'info') } // 验证是否实现了负载均衡 const counts = Object.values(selectionCount) const avgCount = totalSelections / accounts.length const variance = counts.reduce((sum, count) => sum + Math.pow(count - avgCount, 2), 0) / counts.length const stdDev = Math.sqrt(variance) log(`\n 平均选择次数: ${avgCount.toFixed(1)}`, 'info') log(` 标准差: ${stdDev.toFixed(1)}`, 'info') // 如果标准差小于平均值的50%,认为负载均衡效果良好 if (stdDev < avgCount * 0.5) { log('✅ 负载均衡验证通过,分布相对均匀', 'success') } else { log('⚠️ 负载分布不够均匀,但这可能是正常的随机波动', 'warning') } } catch (error) { log(`❌ 测试5失败: ${error.message}`, 'error') throw error } } // 测试6: 会话粘性测试 async function test6_stickySession() { log('\n📝 测试6: 会话粘性(Sticky Session)测试', 'info') try { const apiKey = testData.apiKeys[0] const sessionHash = `test_session_${uuidv4()}` // 第一次选择 const firstSelection = await unifiedClaudeScheduler.selectAccountForApiKey( { id: apiKey.id, claudeAccountId: apiKey.claudeAccountId, name: apiKey.name }, sessionHash ) log(` 首次选择账户: ${firstSelection.accountId}`, 'info') // 使用相同的sessionHash多次请求 let consistentCount = 0 const testCount = 10 for (let i = 0; i < testCount; i++) { const selection = await unifiedClaudeScheduler.selectAccountForApiKey( { id: apiKey.id, claudeAccountId: apiKey.claudeAccountId, name: apiKey.name }, sessionHash ) if (selection.accountId === firstSelection.accountId) { consistentCount++ } await sleep(100) } log(` 会话一致性: ${consistentCount}/${testCount} 次选择了相同账户`, 'info') if (consistentCount === testCount) { log('✅ 会话粘性验证通过,同一会话始终选择相同账户', 'success') } else { throw new Error(`会话粘性失败,只有${consistentCount}/${testCount}次选择了相同账户`) } } catch (error) { log(`❌ 测试6失败: ${error.message}`, 'error') throw error } } // 测试7: 账户可用性检查 async function test7_accountAvailability() { log('\n📝 测试7: 账户可用性检查', 'info') try { const apiKey = testData.apiKeys[0] const accounts = testData.accounts.filter( (a) => a.type === 'claude' || a.type === 'claude-console' ) // 禁用第一个账户 const firstAccount = accounts[0] if (firstAccount.type === 'claude') { await claudeAccountService.updateAccount(firstAccount.id, { isActive: false }) } else { await claudeConsoleAccountService.updateAccount(firstAccount.id, { isActive: false }) } log(` 已禁用账户: ${firstAccount.name}`, 'info') // 多次选择,验证不会选择到禁用的账户 const selectionResults = [] for (let i = 0; i < 20; i++) { const sessionHash = uuidv4() // 每次使用新会话 const result = await unifiedClaudeScheduler.selectAccountForApiKey( { id: apiKey.id, claudeAccountId: apiKey.claudeAccountId, name: apiKey.name }, sessionHash ) selectionResults.push(result.accountId) } // 检查是否选择了禁用的账户 const selectedDisabled = selectionResults.includes(firstAccount.id) if (!selectedDisabled) { log('✅ 账户可用性验证通过,未选择禁用的账户', 'success') } else { throw new Error('错误:选择了已禁用的账户') } // 重新启用账户 if (firstAccount.type === 'claude') { await claudeAccountService.updateAccount(firstAccount.id, { isActive: true }) } else { await claudeConsoleAccountService.updateAccount(firstAccount.id, { isActive: true }) } } catch (error) { log(`❌ 测试7失败: ${error.message}`, 'error') throw error } } // 测试8: 分组成员管理 async function test8_groupMemberManagement() { log('\n📝 测试8: 分组成员管理', 'info') try { const claudeGroup = testData.groups.find((g) => g.platform === 'claude') const account = testData.accounts.find((a) => a.type === 'claude') // 获取账户所属分组 const accountGroups = await accountGroupService.getAccountGroup(account.id) const hasTargetGroup = accountGroups.some((group) => group.id === claudeGroup.id) if (hasTargetGroup) { log('✅ 账户分组查询验证通过', 'success') } else { throw new Error('账户分组查询结果不正确') } // 从分组移除账户 await accountGroupService.removeAccountFromGroup(account.id, claudeGroup.id) log(` 从分组移除账户: ${account.name}`, 'info') // 验证账户已不在分组中 const membersAfterRemove = await accountGroupService.getGroupMembers(claudeGroup.id) if (!membersAfterRemove.includes(account.id)) { log('✅ 账户移除验证通过', 'success') } else { throw new Error('账户移除失败') } // 重新添加账户 await accountGroupService.addAccountToGroup(account.id, claudeGroup.id, 'claude') log(' 重新添加账户到分组', 'info') } catch (error) { log(`❌ 测试8失败: ${error.message}`, 'error') throw error } } // 测试9: 空分组处理 async function test9_emptyGroupHandling() { log('\n📝 测试9: 空分组处理', 'info') try { // 创建一个空分组 const emptyGroup = await accountGroupService.createGroup({ name: `${TEST_PREFIX}空分组`, platform: 'claude', description: '测试空分组' }) testData.groups.push(emptyGroup) // 创建绑定到空分组的API Key const apiKey = await apiKeyService.generateApiKey({ name: `${TEST_PREFIX}空分组API Key`, claudeAccountId: `group:${emptyGroup.id}`, permissions: 'claude' }) testData.apiKeys.push(apiKey) // 尝试从空分组选择账户(应该失败) try { await unifiedClaudeScheduler.selectAccountForApiKey({ id: apiKey.id, claudeAccountId: apiKey.claudeAccountId, name: apiKey.name }) throw new Error('空分组选择账户应该失败') } catch (error) { if (error.message.includes('has no members')) { log(`✅ 空分组处理验证通过:${error.message}`, 'success') } else { throw error } } } catch (error) { log(`❌ 测试9失败: ${error.message}`, 'error') throw error } } // 主测试函数 async function runTests() { log('\n🚀 开始分组调度功能测试\n', 'info') try { // 连接Redis await redis.connect() log('✅ Redis连接成功', 'success') // 执行测试 await test1_createGroups() await test2_createAccountsAndAddToGroup() await test3_platformConsistency() await test4_apiKeyBindGroup() await test5_groupSchedulingLoadBalance() await test6_stickySession() await test7_accountAvailability() await test8_groupMemberManagement() await test9_emptyGroupHandling() log('\n🎉 所有测试通过!分组调度功能工作正常', 'success') } catch (error) { log(`\n❌ 测试失败: ${error.message}`, 'error') console.error(error) } finally { // 清理测试数据 if (CLEANUP_ON_FINISH) { await cleanup() } else { log('\n⚠️ 测试数据未清理,请手动清理', 'warning') } // 关闭Redis连接 await redis.disconnect() process.exit(0) } } // 运行测试 runTests()