feat: add grouping for mock data (#4151)

* feat: add grouping for mock data

* docs: update api

packages/core/src/service/mockService.ts

@@ -13,77 +13,116 @@ import {
   Scope,
   MOCK_KEY,
 } from '../decorator';
+import { isClass } from '../util/types';
 
 @Provide()
 @Scope(ScopeEnum.Singleton)
 export class MidwayMockService {
-  protected mocks = [];
-  protected contextMocks: Array<{
-    app: IMidwayApplication;
-    key: string | ((ctx: IMidwayContext) => void);
-    value: any;
-  }> = [];
-  protected cache = new Map();
+  /**
+   * Save class prototype and object property mocks
+   */
+  protected mocks: Map<
+    string,
+    Array<{
+      obj: any;
+      key: string;
+      descriptor: PropertyDescriptor;
+      hasOwnProperty: boolean;
+    }>
+  > = new Map();
+  /**
+   * Save context mocks
+   */
+  protected contextMocks: Map<
+    string,
+    Array<{
+      app: IMidwayApplication;
+      key: string | ((ctx: IMidwayContext) => void);
+      value: any;
+    }>
+  > = new Map();
+  protected cache: Map<string, Map<any, Set<string>>> = new Map();
   protected simulatorList: Array<ISimulation> = [];
   constructor(readonly applicationContext: IMidwayContainer) {}
 
   @Init()
   async init() {
     if (MidwayMockService.prepareMocks.length > 0) {
       for (const item of MidwayMockService.prepareMocks) {
-        this.mockProperty(item.obj, item.key, item.value);
+        this.mockProperty(item.obj, item.key, item.value, item.group);
       }
       MidwayMockService.prepareMocks = [];
     }
   }
 
+  /**
+   * Prepare mocks before the service is initialized
+   */
   static prepareMocks = [];
 
   static mockClassProperty(
     clzz: new (...args) => any,
     propertyName: string,
-    value: any
+    value: any,
+    group = 'default'
   ) {
-    this.mockProperty(clzz.prototype, propertyName, value);
+    this.mockProperty(clzz.prototype, propertyName, value, group);
   }
 
-  static mockProperty(obj: new (...args) => any, key: string, value: any) {
+  static mockProperty(
+    obj: new (...args) => any,
+    key: string,
+    value: any,
+    group = 'default'
+  ) {
     this.prepareMocks.push({
       obj,
       key,
       value,
+      group,
     });
   }
 
   public mockClassProperty(
     clzz: new (...args) => any,
     propertyName: string,
-    value: any
+    value: any,
+    group = 'default'
   ) {
-    return this.mockProperty(clzz.prototype, propertyName, value);
+    return this.mockProperty(clzz.prototype, propertyName, value, group);
   }
 
-  public mockProperty(obj: any, key: string, value) {
+  public mockProperty(obj: any, key: string, value: any, group = 'default') {
     // eslint-disable-next-line no-prototype-builtins
     const hasOwnProperty = obj.hasOwnProperty(key);
-    this.mocks.push({
+    const mockItem = {
       obj,
       key,
       descriptor: Object.getOwnPropertyDescriptor(obj, key),
       // Make sure the key exists on object not the prototype
       hasOwnProperty,
-    });
+    };
+
+    if (!this.mocks.has(group)) {
+      this.mocks.set(group, []);
+    }
+    this.mocks.get(group).push(mockItem);
 
-    // Delete the origin key, redefine it below
     if (hasOwnProperty) {
       delete obj[key];
     }
 
     // Set a flag that checks if it is mocked
-    let flag = this.cache.get(obj);
+    let groupCache = this.cache.get(group);
+    if (!groupCache) {
+      groupCache = new Map();
+      this.cache.set(group, groupCache);
+    }
+
+    let flag = groupCache.get(obj);
     if (!flag) {
       flag = new Set();
-      this.cache.set(obj, flag);
+      groupCache.set(obj, flag);
     }
     flag.add(key);
 
@@ -94,42 +133,68 @@ export class MidwayMockService {
   public mockContext(
     app: IMidwayApplication,
     key: string | ((ctx: IMidwayContext) => void),
-    value?: PropertyDescriptor | any
+    value?: PropertyDescriptor | any,
+    group = 'default'
   ) {
-    this.contextMocks.push({
+    if (!this.contextMocks.has(group)) {
+      this.contextMocks.set(group, []);
+    }
+    this.contextMocks.get(group).push({
       app,
       key,
       value,
     });
   }
 
+  public restore(group = 'default') {
+    this.restoreGroup(group);
+  }
+
   @Destroy()
-  restore() {
-    for (let i = this.mocks.length - 1; i >= 0; i--) {
-      const m = this.mocks[i];
+  public restoreAll() {
+    const groups = new Set([
+      ...this.mocks.keys(),
+      ...this.contextMocks.keys(),
+      ...this.cache.keys(),
+    ]);
+
+    for (const group of groups) {
+      this.restoreGroup(group);
+    }
+
+    this.simulatorList = [];
+  }
+
+  private restoreGroup(group: string) {
+    const groupMocks = this.mocks.get(group) || [];
+    for (let i = groupMocks.length - 1; i >= 0; i--) {
+      const m = groupMocks[i];
       if (!m.hasOwnProperty) {
-        // Delete the mock key, use key on the prototype
         delete m.obj[m.key];
       } else {
-        // Redefine the origin key instead of the mock key
         Object.defineProperty(m.obj, m.key, m.descriptor);
       }
     }
-    this.mocks = [];
-    this.contextMocks = [];
-    this.cache.clear();
-    this.simulatorList = [];
-    MidwayMockService.prepareMocks = [];
+    this.mocks.delete(group);
+    this.contextMocks.delete(group);
+    this.cache.delete(group);
+    this.simulatorList = this.simulatorList.filter(
+      sim => sim['group'] !== group
+    );
   }
 
-  isMocked(obj, key) {
-    const flag = this.cache.get(obj);
+  public isMocked(obj, key, group = 'default') {
+    if (isClass(obj)) {
+      obj = obj.prototype;
+    }
+    const groupCache = this.cache.get(group);
+    const flag = groupCache ? groupCache.get(obj) : undefined;
     return flag ? flag.has(key) : false;
   }
 
   applyContextMocks(app: IMidwayApplication, ctx: IMidwayContext) {
-    if (this.contextMocks.length > 0) {
-      for (const mockItem of this.contextMocks) {
+    for (const [, groupMocks] of this.contextMocks) {
+      for (const mockItem of groupMocks) {
         if (mockItem.app === app) {
           const descriptor = this.overridePropertyDescriptor(mockItem.value);
           if (typeof mockItem.key === 'string') {
@@ -143,7 +208,10 @@ export class MidwayMockService {
   }
 
   getContextMocksSize() {
-    return this.contextMocks.length;
+    return Array.from(this.contextMocks.values()).reduce(
+      (sum, group) => sum + group.length,
+      0
+    );
   }
 
   private overridePropertyDescriptor(value) {
@@ -165,13 +233,14 @@ export class MidwayMockService {
     return descriptor;
   }
 
-  public async initSimulation() {
+  public async initSimulation(group = 'default') {
     const simulationModule: Array<new (...args) => ISimulation> =
       listModule(MOCK_KEY);
 
     for (const module of simulationModule) {
       const instance = await this.applicationContext.getAsync(module);
       if (await instance.enableCondition()) {
+        instance['group'] = group;
         this.simulatorList.push(instance);
       }
     }

packages/core/test/service/mockService.test.ts

@@ -4,15 +4,24 @@ import { MidwayMockService, IMidwayApplication } from '../../src';
 import { UserService } from '../fixtures/base-app-ctx-mock/src/configuration';
 
 describe('/service/mockService.test.ts', () => {
+  let framework;
+  let app: IMidwayApplication;
+  let mockService: MidwayMockService;
 
-  it('should test mock context', async () => {
-    const framework = await createLightFramework(path.join(
+  beforeAll(async () => {
+    framework = await createLightFramework(path.join(
       __dirname,
       './fixtures/base-app-ctx-mock/src'
     ));
+    app = framework.getApplication() as IMidwayApplication;
+    mockService = framework.getApplicationContext().get(MidwayMockService);
+  });
 
-    const app = framework.getApplication() as IMidwayApplication;
-    const mockService = framework.getApplicationContext().get(MidwayMockService);
+  afterAll(async () => {
+    await framework.stop();
+  });
+
+  it('should test mock context', async () => {
     mockService.mockContext(app, 'user', 'zhangting');
     mockService.mockContext(app, (ctx) => {
       ctx['bbbb'] = 'cccc';
@@ -62,7 +71,116 @@ describe('/service/mockService.test.ts', () => {
     mockService.restore();
 
     expect(await userService.invoke()).toEqual('hello world');
+  });
 
-    await framework.stop();
+  it('should test mock with groups', async () => {
+    // 测试不同分组的 mock
+    mockService.mockContext(app, 'user', 'zhangting', 'group1');
+    mockService.mockContext(app, 'role', 'admin', 'group2');
+
+    let ctx = app.createAnonymousContext();
+    const fn = await framework.applyMiddleware();
+    await fn((ctx));
+
+    expect(ctx['user']).toEqual('zhangting');
+    expect(ctx['role']).toEqual('admin');
+
+    // 测试恢复单个分组
+    mockService.restore('group1');
+
+    ctx = app.createAnonymousContext();
+    await (await framework.applyMiddleware())(ctx);
+
+    expect(ctx['user']).toBeUndefined();
+    expect(ctx['role']).toEqual('admin');
+
+    // 测试恢复所有分组
+    mockService.restoreAll();
+
+    ctx = app.createAnonymousContext();
+    await (await framework.applyMiddleware())(ctx);
+
+    expect(ctx['user']).toBeUndefined();
+    expect(ctx['role']).toBeUndefined();
+  });
+
+  it('should test mock class property with groups', async () => {
+    framework.getApplicationContext().bindClass(UserService);
+    mockService.mockClassProperty(UserService, 'invoke', () => '1112', 'group1');
+    mockService.mockClassProperty(UserService, 'getName', () => 'mockName', 'group2');
+
+    expect(mockService.isMocked(UserService, 'invoke', 'group1')).toBeTruthy();
+
+    const userService = await framework.getApplicationContext().getAsync(UserService);
+    expect(userService.invoke()).toEqual('1112');
+    expect(userService.getName()).toEqual('mockName');
+
+    // 恢复单个分组
+    mockService.restore('group1');
+
+    expect(await userService.invoke()).toEqual('hello world');
+    expect(userService.getName()).toEqual('mockName');
+
+    // 恢复所有分组
+    mockService.restoreAll();
+
+    expect(mockService.isMocked(UserService, 'invoke', 'group1')).toBeFalsy();
+    expect(await userService.invoke()).toEqual('hello world');
+    expect(() => {
+      userService.getName();
+    }).toThrow('userService.getName is not a function');
+  });
+
+  it('should test mock property with groups', async () => {
+    const obj = {
+      method1: () => 'original1',
+      method2: () => 'original2'
+    };
+
+    mockService.mockProperty(obj, 'method1', () => 'mocked1', 'group1');
+    mockService.mockProperty(obj, 'method2', () => 'mocked2', 'group2');
+
+    expect(obj.method1()).toEqual('mocked1');
+    expect(obj.method2()).toEqual('mocked2');
+
+    // 恢复单个分组
+    mockService.restore('group1');
+
+    expect(obj.method1()).toEqual('original1');
+    expect(obj.method2()).toEqual('mocked2');
+
+    // 恢复所有分组
+    mockService.restoreAll();
+
+    expect(obj.method1()).toEqual('original1');
+    expect(obj.method2()).toEqual('original2');
+  });
+
+  it('should test isMocked with groups', async () => {
+    const obj = { method: () => 'original' };
+
+    mockService.mockProperty(obj, 'method', () => 'mocked', 'testGroup');
+
+    expect(mockService.isMocked(obj, 'method', 'testGroup')).toBeTruthy();
+
+    mockService.restore('testGroup');
+
+    expect(mockService.isMocked(obj, 'method', 'testGroup')).toBeFalsy();
+  });
+
+  it('should test mock without specifying group', async () => {
+    const obj = { method: () => 'original' };
+
+    // 不传 group，使用默认分组
+    mockService.mockProperty(obj, 'method', () => 'mocked');
+
+    expect(mockService.isMocked(obj, 'method')).toBeTruthy();
+    expect(obj.method()).toEqual('mocked');
+
+    // 恢复默认分组
+    mockService.restore();
+
+    expect(mockService.isMocked(obj, 'method')).toBeFalsy();
+    expect(obj.method()).toEqual('original');
   });
 });