claude-relay-service/scripts/test-group-scheduling.js

/**
 * 分组调度功能测试脚本
 * 用于测试账户分组管理和调度逻辑的正确性
 */

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