file-starter/tests/utils/server/test-db.ts

import { DataSource, EntityManager, Repository } from 'typeorm';
import { vi, beforeEach, afterEach } from 'vitest';

/**
 * 创建模拟的数据源
 */
export function createMockDataSource() {
  const manager = createMockEntityManager();
  const mockDataSource = {
    initialize: vi.fn().mockResolvedValue(undefined),
    destroy: vi.fn().mockResolvedValue(undefined),
    isInitialized: true,
    manager,
    getRepository: vi.fn().mockImplementation(() => createMockRepository()),
    createQueryBuilder: vi.fn().mockReturnValue(createMockQueryBuilder()),
    transaction: vi.fn().mockImplementation(async (callback) => {
      return callback(manager);
    }),
    synchronize: vi.fn().mockResolvedValue(undefined),
    dropDatabase: vi.fn().mockResolvedValue(undefined)
  };

  return mockDataSource;
}

/**
 * 创建模拟的实体管理器
 */
export function createMockEntityManager(): EntityManager {
  return {
    find: vi.fn().mockResolvedValue([]),
    findOne: vi.fn().mockResolvedValue(null),
    save: vi.fn().mockImplementation((entity) => Promise.resolve(entity)),
    update: vi.fn().mockResolvedValue({ affected: 1 }),
    delete: vi.fn().mockResolvedValue({ affected: 1 }),
    createQueryBuilder: vi.fn().mockReturnValue(createMockQueryBuilder()),
    transaction: vi.fn().mockImplementation(async (callback) => {
      return callback(createMockEntityManager());
    }),
    getRepository: vi.fn().mockImplementation(() => createMockRepository())
  } as any;
}

/**
 * 创建模拟的Repository
 */
export function createMockRepository<T extends object = any>(): Repository<T> {
  return {
    find: vi.fn().mockResolvedValue([]),
    findOne: vi.fn().mockResolvedValue(null),
    findOneBy: vi.fn().mockResolvedValue(null),
    findOneByOrFail: vi.fn().mockResolvedValue(null),
    findBy: vi.fn().mockResolvedValue([]),
    findAndCount: vi.fn().mockResolvedValue([[], 0]),
    findAndCountBy: vi.fn().mockResolvedValue([[], 0]),
    save: vi.fn().mockImplementation((entity) => Promise.resolve(entity)),
    update: vi.fn().mockResolvedValue({ affected: 1 }),
    delete: vi.fn().mockResolvedValue({ affected: 1 }),
    create: vi.fn().mockImplementation((entity) => ({ ...entity, id: Date.now() })),
    createQueryBuilder: vi.fn().mockReturnValue(createMockQueryBuilder()),
    count: vi.fn().mockResolvedValue(0),
    countBy: vi.fn().mockResolvedValue(0),
    exist: vi.fn().mockResolvedValue(false)
  } as any;
}

/**
 * 创建模拟的QueryBuilder
 */
export function createMockQueryBuilder() {
  const mockQueryBuilder = {
    select: vi.fn().mockReturnThis(),
    from: vi.fn().mockReturnThis(),
    where: vi.fn().mockReturnThis(),
    andWhere: vi.fn().mockReturnThis(),
    orWhere: vi.fn().mockReturnThis(),
    leftJoin: vi.fn().mockReturnThis(),
    innerJoin: vi.fn().mockReturnThis(),
    orderBy: vi.fn().mockReturnThis(),
    groupBy: vi.fn().mockReturnThis(),
    having: vi.fn().mockReturnThis(),
    skip: vi.fn().mockReturnThis(),
    take: vi.fn().mockReturnThis(),
    getMany: vi.fn().mockResolvedValue([]),
    getOne: vi.fn().mockResolvedValue(null),
    getCount: vi.fn().mockResolvedValue(0),
    getRawMany: vi.fn().mockResolvedValue([]),
    getRawOne: vi.fn().mockResolvedValue(null),
    execute: vi.fn().mockResolvedValue(undefined),
    setParameter: vi.fn().mockReturnThis(),
    setParameters: vi.fn().mockReturnThis()
  };

  return mockQueryBuilder;
}

/**
 * 数据库测试工具类
 */
export class TestDatabase {
  private static dataSource: DataSource | null = null;

  /**
   * 初始化测试数据库
   */
  static async initialize(): Promise<DataSource> {
    if (this.dataSource?.isInitialized) {
      return this.dataSource;
    }

    // 使用SQLite内存数据库进行测试
    this.dataSource = new DataSource({
      type: 'better-sqlite3',
      database: ':memory:',
      synchronize: true,
      logging: false,
      entities: [
        // 导入实际实体
        (await import('@/server/modules/users/user.entity')).UserEntity,
        (await import('@/server/modules/users/role.entity')).Role
      ]
    });

    await this.dataSource.initialize();
    return this.dataSource;
  }

  /**
   * 清理测试数据库
   */
  static async cleanup(): Promise<void> {
    if (this.dataSource?.isInitialized) {
      await this.dataSource.destroy();
      this.dataSource = null;
    }
  }

  /**
   * 获取当前数据源
   */
  static getDataSource(): DataSource | null {
    return this.dataSource;
  }
}

/**
 * 测试数据库生命周期钩子
 */
export function setupDatabaseHooks() {
  beforeEach(async () => {
    await TestDatabase.initialize();
  });

  afterEach(async () => {
    await TestDatabase.cleanup();
  });
}