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(),
    setSessionManager: vi.fn(),
    setAgentMode: vi.fn(),
    chat: vi.fn().mockResolvedValue({
      text: 'mock response',
      messages: [
        { role: 'user', content: 'test' },
        { role: 'assistant', content: '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),
    }),
    getHistory: vi.fn().mockReturnValue([
      { role: 'user', content: 'test' },
      { role: 'assistant', content: 'mock response' },
    ]),
  };
}

/**
 * 创建 Mock SessionManager 实例
 */
export function createMockSessionManager() {
  return {
    init: vi.fn().mockResolvedValue(undefined),
    getSession: vi.fn().mockReturnValue({ id: 'session-1', messages: [] }),
    setMessages: vi.fn().mockResolvedValue(undefined),
    setDiscoveredTools: vi.fn().mockResolvedValue(undefined),
    save: vi.fn().mockResolvedValue(undefined),
    close: vi.fn().mockResolvedValue(undefined),
    restoreSession: vi.fn().mockResolvedValue({ id: 'session-1', messages: [] }),
  };
}

/**
 * 创建 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,
  };
}

/**
 * 创建 Mock AgentEventEmitter
 */
export function createMockAgentEventEmitter() {
  return {
    on: vi.fn().mockReturnValue(() => {}), // 返回 unsubscribe 函数
    off: vi.fn(),
    emit: vi.fn(),
    clear: vi.fn(),
  };
}

// 全局 mock 状态，用于测试时配置
let mockProviderRegistry = createMockProviderRegistry();
let mockAgentRegistry = createMockAgentRegistry();
let mockPermissionManager = createMockPermissionManager();
let mockToolRegistry = createMockToolRegistry();
let mockAgentEventEmitter = createMockAgentEventEmitter();
let mockAgent = createMockAgent();
let mockSessionManager = createMockSessionManager();
let loadConfigFn = vi.fn().mockReturnValue({
  provider: 'anthropic',
  apiKey: 'test-api-key',
  model: 'claude-sonnet-4-20250514',
  maxTokens: 4096,
  systemPrompt: 'test prompt',
});

/**
 * 重置所有 mock 到默认状态
 */
export function resetCoreMocks() {
  mockProviderRegistry = createMockProviderRegistry();
  mockAgentRegistry = createMockAgentRegistry();
  mockPermissionManager = createMockPermissionManager();
  mockToolRegistry = createMockToolRegistry();
  mockAgentEventEmitter = createMockAgentEventEmitter();
  mockAgent = createMockAgent();
  mockSessionManager = createMockSessionManager();
  loadConfigFn = vi.fn().mockReturnValue({
    provider: 'anthropic',
    apiKey: 'test-api-key',
    model: 'claude-sonnet-4-20250514',
    maxTokens: 4096,
    systemPrompt: 'test prompt',
  });
}

/**
 * 配置 mock 行为
 */
export function configureMocks(config: {
  providerRegistry?: Partial<ReturnType<typeof createMockProviderRegistry>>;
  agentRegistry?: Partial<ReturnType<typeof createMockAgentRegistry>>;
  permissionManager?: Partial<ReturnType<typeof createMockPermissionManager>>;
  toolRegistry?: Partial<ReturnType<typeof createMockToolRegistry>>;
  agentEventEmitter?: Partial<ReturnType<typeof createMockAgentEventEmitter>>;
  agent?: Partial<ReturnType<typeof createMockAgent>>;
  sessionManager?: Partial<ReturnType<typeof createMockSessionManager>>;
  loadConfig?: ReturnType<typeof vi.fn>;
}) {
  if (config.providerRegistry) {
    Object.assign(mockProviderRegistry, config.providerRegistry);
  }
  if (config.agentRegistry) {
    Object.assign(mockAgentRegistry, config.agentRegistry);
  }
  if (config.permissionManager) {
    Object.assign(mockPermissionManager, config.permissionManager);
  }
  if (config.toolRegistry) {
    Object.assign(mockToolRegistry, config.toolRegistry);
  }
  if (config.agentEventEmitter) {
    Object.assign(mockAgentEventEmitter, config.agentEventEmitter);
  }
  if (config.agent) {
    Object.assign(mockAgent, config.agent);
  }
  if (config.sessionManager) {
    Object.assign(mockSessionManager, config.sessionManager);
  }
  if (config.loadConfig) {
    loadConfigFn = config.loadConfig;
  }
}

/**
 * 获取当前 mock 状态（用于测试验证）
 */
export function getMocks() {
  return {
    providerRegistry: mockProviderRegistry,
    agentRegistry: mockAgentRegistry,
    permissionManager: mockPermissionManager,
    toolRegistry: mockToolRegistry,
    agentEventEmitter: mockAgentEventEmitter,
    agent: mockAgent,
    sessionManager: mockSessionManager,
    loadConfig: loadConfigFn,
  };
}

// Mock Agent 类
export const MockAgent = vi.fn().mockImplementation(function (this: any) {
  Object.assign(this, mockAgent);
  return this;
});

// Mock SessionManager 类
export const MockSessionManager = vi.fn().mockImplementation(function (this: any) {
  Object.assign(this, mockSessionManager);
  return this;
});

// Mock ConfigurationError 类
export class MockConfigurationError extends Error {
  provider: string;

  constructor(message: string, provider: string = 'anthropic') {
    super(message);
    this.name = 'ConfigurationError';
    this.provider = provider;
  }
}

/**
 * 导出完整的 Core 模块 mock
 * 用于 vi.mock('@ai-assistant/core', ...)
 */
export const coreMock = {
  Agent: MockAgent,
  SessionManager: MockSessionManager,
  toolRegistry: mockToolRegistry,
  loadConfig: () => loadConfigFn(),
  getPermissionManager: () => mockPermissionManager,
  getProviderRegistry: () => mockProviderRegistry,
  agentRegistry: mockAgentRegistry,
  agentEventEmitter: mockAgentEventEmitter,
  ConfigurationError: MockConfigurationError,
};

// 兼容旧 API
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>;
} = {}) {
  configureMocks({
    agent: overrides.agent,
    providerRegistry: overrides.providerRegistry,
    agentRegistry: overrides.agentRegistry,
    permissionManager: overrides.permissionManager,
    toolRegistry: overrides.toolRegistry,
    loadConfig: overrides.loadConfig,
  });

  return coreMock;
}

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