ai-terminal-assistant/packages/server/tests/mocks/core.mock.ts

/**
 * Core 模块 Mock 工厂
 *
 * 提供 @ai-assistant/core 模块的 mock 实现
 */

import { vi } from 'vitest';

/**
 * 创建 Mock Agent 实例
 */
export function createMockAgent() {
  return {
    setRegistry: vi.fn(),
    chat: vi.fn().mockResolvedValue('mock response'),
    getToolCount: vi.fn().mockReturnValue({ core: 5, discovered: 0, total: 5 }),
    getContextUsageFormatted: vi.fn().mockReturnValue('10k / 200k'),
    getContextUsage: vi.fn().mockReturnValue({
      input: 10000,
      contextLimit: 200000,
      available: 190000,
      usagePercent: 5,
    }),
    compactHistory: vi.fn().mockResolvedValue({ freedTokens: 1000, type: 'prune' }),
    getCompressionManager: vi.fn().mockReturnValue({
      shouldCompress: vi.fn().mockReturnValue(false),
    }),
  };
}

/**
 * 创建 Mock ProviderRegistry
 */
export function createMockProviderRegistry(overrides: Partial<ReturnType<typeof createMockProviderRegistry>> = {}) {
  return {
    init: vi.fn().mockResolvedValue(undefined),
    isInitialized: vi.fn().mockReturnValue(false),
    list: vi.fn().mockReturnValue([]),
    listForApi: vi.fn().mockReturnValue([]),
    get: vi.fn(),
    getInfo: vi.fn(),
    getDetail: vi.fn(),
    has: vi.fn().mockReturnValue(false),
    getConfig: vi.fn(),
    setConfig: vi.fn(),
    getAllConfigs: vi.fn().mockReturnValue({}),
    getModels: vi.fn().mockReturnValue([]),
    getModelInfo: vi.fn(),
    addCustomModel: vi.fn(),
    removeCustomModel: vi.fn().mockReturnValue(false),
    testConnection: vi.fn().mockResolvedValue({ success: true }),
    getModelFactory: vi.fn().mockReturnValue(() => ({})),
    saveConfig: vi.fn().mockResolvedValue(undefined),
    reloadConfig: vi.fn().mockResolvedValue(undefined),
    ...overrides,
  };
}

/**
 * 创建 Mock AgentRegistry
 */
export function createMockAgentRegistry(overrides: Partial<ReturnType<typeof createMockAgentRegistry>> = {}) {
  return {
    init: vi.fn().mockResolvedValue(undefined),
    isInitialized: vi.fn().mockReturnValue(false),
    get: vi.fn(),
    list: vi.fn().mockReturnValue([]),
    listSubagents: vi.fn().mockReturnValue([]),
    listPrimaryAgents: vi.fn().mockReturnValue([]),
    getInternal: vi.fn(),
    listInternalAgents: vi.fn().mockReturnValue([]),
    register: vi.fn(),
    remove: vi.fn().mockReturnValue(false),
    has: vi.fn().mockReturnValue(false),
    size: 0,
    getNames: vi.fn().mockReturnValue([]),
    getGlobalConfig: vi.fn().mockReturnValue(null),
    generateSubagentDescription: vi.fn().mockReturnValue(''),
    ...overrides,
  };
}

/**
 * 创建 Mock PermissionManager
 */
export function createMockPermissionManager(overrides: Partial<ReturnType<typeof createMockPermissionManager>> = {}) {
  return {
    setAskCallback: vi.fn(),
    checkPermission: vi.fn().mockResolvedValue({ allowed: true }),
    checkFilePermission: vi.fn().mockResolvedValue({ allowed: true }),
    checkBashPermission: vi.fn().mockResolvedValue({ allowed: true }),
    ...overrides,
  };
}

/**
 * 创建 Mock ToolRegistry
 */
export function createMockToolRegistry(overrides: Partial<ReturnType<typeof createMockToolRegistry>> = {}) {
  return {
    getCoreTools: vi.fn().mockReturnValue([]),
    getAllTools: vi.fn().mockReturnValue([]),
    get: vi.fn(),
    has: vi.fn().mockReturnValue(false),
    register: vi.fn(),
    ...overrides,
  };
}

/**
 * 创建完整的 Core 模块 mock
 */
export function createMockCoreModule(overrides: {
  agent?: Partial<ReturnType<typeof createMockAgent>>;
  providerRegistry?: Partial<ReturnType<typeof createMockProviderRegistry>>;
  agentRegistry?: Partial<ReturnType<typeof createMockAgentRegistry>>;
  permissionManager?: Partial<ReturnType<typeof createMockPermissionManager>>;
  toolRegistry?: Partial<ReturnType<typeof createMockToolRegistry>>;
  loadConfig?: ReturnType<typeof vi.fn>;
  saveConfig?: ReturnType<typeof vi.fn>;
} = {}) {
  const mockAgent = { ...createMockAgent(), ...overrides.agent };
  const mockProviderRegistry = createMockProviderRegistry(overrides.providerRegistry);
  const mockAgentRegistry = createMockAgentRegistry(overrides.agentRegistry);
  const mockPermissionManager = createMockPermissionManager(overrides.permissionManager);
  const mockToolRegistry = createMockToolRegistry(overrides.toolRegistry);

  // 创建一个真正的类来模拟 Agent 构造函数
  const MockAgentClass = vi.fn().mockImplementation(function (this: any) {
    Object.assign(this, mockAgent);
    return this;
  });

  return {
    Agent: MockAgentClass,
    toolRegistry: mockToolRegistry,
    loadConfig: overrides.loadConfig ?? vi.fn().mockReturnValue({
      provider: 'anthropic',
      apiKey: 'test-api-key',
      model: 'claude-sonnet-4-20250514',
      maxTokens: 4096,
      systemPrompt: 'test prompt',
    }),
    saveConfig: overrides.saveConfig ?? vi.fn(),
    getPermissionManager: vi.fn().mockReturnValue(mockPermissionManager),
    getProviderRegistry: vi.fn().mockReturnValue(mockProviderRegistry),
    agentRegistry: mockAgentRegistry,
    // 额外暴露内部 mock 以便测试验证
    _mockAgent: mockAgent,
    _mockProviderRegistry: mockProviderRegistry,
    _mockAgentRegistry: mockAgentRegistry,
    _mockPermissionManager: mockPermissionManager,
    _mockToolRegistry: mockToolRegistry,
  };
}

/**
 * 创建抛出 ConfigurationError 的 loadConfig mock
 */
export function createConfigurationErrorMock(message: string, provider: string = 'anthropic') {
  return vi.fn().mockImplementation(() => {
    const error = new Error(message) as Error & { provider: string };
    error.name = 'ConfigurationError';
    error.provider = provider;
    throw error;
  });
}