replace global mockInstrument function with a local one (#66)

XMadrid · web-flow · commit 1dc6d9b2398d · 2025-09-23T15:39:23.000+08:00
diff --git a/assembly/implement.ts b/assembly/implement.ts
@@ -16,19 +16,22 @@ export function testImpl(name: string, testFunction: () => void): void {
 }
 
 export function mockImpl<T extends Function>(
-  oldFunction: T,
-  newFunction: T,
+  originalFunction: T,
+  mockFunction: T,
 ): MockFn {
-  if (!isFunction<T>(oldFunction) || !isFunction<T>(newFunction)) {
+  if (!isFunction<T>(originalFunction) || !isFunction<T>(mockFunction)) {
     ERROR("mock paramemter receive a function");
   }
-  const mockFn = new MockFn(oldFunction.index, newFunction.index);
-  mockFunctionStatus.set(oldFunction.index, newFunction.index);
+  const mockFn = new MockFn(originalFunction.index, mockFunction.index);
+  mockFunctionStatus.setMockFunction(
+    originalFunction.index,
+    mockFunction.index,
+  );
   return mockFn;
 }
-export function unmockImpl<T extends Function>(oldFunction: T): void {
-  mockFunctionStatus.setIgnore(oldFunction.index, true);
+export function unmockImpl<T extends Function>(originalFunction: T): void {
+  mockFunctionStatus.setMockedFunctionIgnore(originalFunction.index, true);
 }
-export function remockImpl<T extends Function>(oldFunction: T): void {
-  mockFunctionStatus.setIgnore(oldFunction.index, false);
+export function remockImpl<T extends Function>(originalFunction: T): void {
+  mockFunctionStatus.setMockedFunctionIgnore(originalFunction.index, false);
 }
diff --git a/assembly/index.ts b/assembly/index.ts
@@ -29,27 +29,27 @@ export function test(name: string, testFunction: () => void): void {
 
 /**
  *  mock some function
- * @param oldFunction function you want to mock
- * @param newFunction the new function.
+ * @param originalFunction function you want to mock
+ * @param mockFunction the new function.
  * @returns Mock Status { callTime : u32}
  */
 export function mock<T extends Function>(
-  oldFunction: T,
-  newFunction: T,
+  originalFunction: T,
+  mockFunction: T,
 ): MockFn {
-  return mockImpl<T>(oldFunction, newFunction);
+  return mockImpl<T>(originalFunction, mockFunction);
 }
 /**
  * unmock this function, can only be used in mocked function
  */
-export function unmock<T extends Function>(oldFunction: T): void {
-  unmockImpl(oldFunction);
+export function unmock<T extends Function>(originalFunction: T): void {
+  unmockImpl(originalFunction);
 }
 /**
  * remock this function, can only be used in mocked function. Pair of {unmock}
  */
-export function remock<T extends Function>(oldFunction: T): void {
-  remockImpl(oldFunction);
+export function remock<T extends Function>(originalFunction: T): void {
+  remockImpl(originalFunction);
 }
 
 export function expect<T>(value: T): Value<T> {
diff --git a/assembly/mockInstrument.ts b/assembly/mockInstrument.ts
@@ -1,19 +1,35 @@
-export declare namespace mockFunctionStatus {
-  function clear(): void;
-  function set(k: u32, v: u32): void;
-  function get(k: u32): u32;
-  function lastGet(): u32;
-  function has(k: u32): bool;
-  function getCalls(oldIndex: u32, newIndex: u32): u32;
-  function setIgnore(k: u32, v: bool): void;
+export namespace mockFunctionStatus {
+
+  @external("__unittest_framework_env","setMockFunction")
+  export declare function setMockFunction(
+    originalFunctionIndex: u32,
+    mockFunctionIndex: u32,
+  ): void;
+
+
+  @external("__unittest_framework_env","getMockedFunctionCalls")
+  export declare function getMockedFunctionCalls(
+    originalFunctionIndex: u32,
+    mockFunctionIndex: u32,
+  ): u32;
+
+
+  @external("__unittest_framework_env","setMockedFunctionIgnore")
+  export declare function setMockedFunctionIgnore(
+    originalFunctionIndex: u32,
+    ignore: bool,
+  ): void;
 }
 
 export class MockFn {
   get calls(): u32 {
-    return mockFunctionStatus.getCalls(this.oldIndex, this.newIndex);
+    return mockFunctionStatus.getMockedFunctionCalls(
+      this.originalFunctionIndex,
+      this.mockFunctionIndex,
+    );
   }
   constructor(
-    public oldIndex: u32,
-    public newIndex: u32,
+    public originalFunctionIndex: u32,
+    public mockFunctionIndex: u32,
   ) {}
 }
diff --git a/instrumentation/CovInstrumentationWalker.cpp b/instrumentation/CovInstrumentationWalker.cpp
@@ -26,7 +26,7 @@ void CovInstrumentationWalker::introduceReportFun() noexcept {
     std::array<BinaryenType, 3U> iii_{BinaryenTypeInt32(), BinaryenTypeInt32(),
                                       BinaryenTypeInt32()};
     const BinaryenType iii = BinaryenTypeCreate(iii_.data(), iii_.size());
-    BinaryenAddFunctionImport(module, reportFunName, "covInstrument", "traceExpression", iii,
+    BinaryenAddFunctionImport(module, reportFunName, "__unittest_framework_env", "traceExpression", iii,
                               wasm::Type::none);
   }
 }
diff --git a/instrumentation/MockInstrumentationWalker.cpp b/instrumentation/MockInstrumentationWalker.cpp
@@ -83,7 +83,7 @@ bool MockInstrumentationWalker::mockFunctionDuplicateImportedCheck() const noexc
     std::array<BinaryenType, 2U> ii_ =
         std::array<BinaryenType, 2U>{BinaryenTypeInt32(), BinaryenTypeInt32()};
     const BinaryenType ii = BinaryenTypeCreate(ii_.data(), ii_.size());
-    BinaryenAddFunctionImport(module, this->checkMock.data(), "mockInstrument", "checkMock", ii,
+    BinaryenAddFunctionImport(module, this->checkMock.data(), "__unittest_framework_env", "checkMock", ii,
                               BinaryenTypeInt32());
   }
   return checkRepeat;
diff --git a/instrumentation/MockInstrumentationWalker.hpp b/instrumentation/MockInstrumentationWalker.hpp
@@ -27,7 +27,7 @@ class MockInstrumentationWalker final : public wasm::PostWalker<MockInstrumentat
   /// @param _checkMock
   explicit MockInstrumentationWalker(
       wasm::Module *const _module,
-      const std::string _checkMock = "mockInstrument/checkMock") noexcept
+      const std::string _checkMock = "__unittest_framework_env/checkMock") noexcept
       : module(_module), checkMock(_checkMock), moduleBuilder(wasm::Builder(*_module)) {
     for (const auto &elemSegment : _module->elementSegments) {
       if (elemSegment->type.isFunction()) {
diff --git a/src/core/covRecorder.ts b/src/core/covRecorder.ts
@@ -5,21 +5,19 @@ export class CoverageRecorder {
 
   getCollectionFuncSet(): Record<string, unknown> {
     return {
-      covInstrument: {
-        traceExpression: (functionIndex: number, basicBlockIndex: number, type: number): void => {
-          switch (type) {
-            case 1: // call in
-            case 2: {
-              // call out
-              // do not need for now
-              break;
-            }
-            case 0: {
-              this._runtimeTrace.push([functionIndex, basicBlockIndex]);
-              break;
-            }
+      traceExpression: (functionIndex: number, basicBlockIndex: number, type: number): void => {
+        switch (type) {
+          case 1: // call in
+          case 2: {
+            // call out
+            // do not need for now
+            break;
           }
-        },
+          case 0: {
+            this._runtimeTrace.push([functionIndex, basicBlockIndex]);
+            break;
+          }
+        }
       },
     };
   }
diff --git a/src/core/execute.ts b/src/core/execute.ts
@@ -4,10 +4,10 @@ import { ensureDirSync } from "fs-extra";
 import { instantiate, Imports as ASImports } from "@assemblyscript/loader";
 import { ExecutionResultSummary } from "../executionResult.js";
 import { Imports, ImportsArgument, InstrumentResult } from "../interface.js";
-import { mockInstrumentFunc } from "../utils/import.js";
 import { supplyDefaultFunction } from "../utils/index.js";
 import { parseImportFunctionInfo } from "../utils/wasmparser.js";
 import { ExecutionRecorder, ExecutionResult } from "./executionRecorder.js";
+import { MockStatusRecorder } from "./mockStatusRecorder.js";
 import { CoverageRecorder } from "./covRecorder.js";
 import assert from "node:assert";
 import { ExecutionError, handleWebAssemblyError } from "../utils/errorTraceHandler.js";
@@ -31,14 +31,17 @@ async function nodeExecutor(
 
   const executionRecorder = new ExecutionRecorder();
   const coverageRecorder = new CoverageRecorder();
+  const mockStatusRecorder = new MockStatusRecorder();
 
   const importsArg = new ImportsArgument(executionRecorder);
   const userDefinedImportsObject = imports === undefined ? {} : imports!(importsArg);
   const importObject: ASImports = {
     wasi_snapshot_preview1: wasi.wasiImport,
-    ...executionRecorder.getCollectionFuncSet(importsArg),
-    mockInstrument: mockInstrumentFunc,
-    ...coverageRecorder.getCollectionFuncSet(),
+    __unittest_framework_env: {
+      ...executionRecorder.getCollectionFuncSet(importsArg),
+      ...mockStatusRecorder.getMockFuncSet(),
+      ...coverageRecorder.getCollectionFuncSet(),
+    },
     ...userDefinedImportsObject,
   } as ASImports;
   const binaryBuffer = await readFile(instrumentResult.instrumentedWasm);
@@ -90,7 +93,7 @@ async function nodeExecutor(
         await exceptionHandler(error);
       }
       executionRecorder.finishTestFunction();
-      mockInstrumentFunc["mockFunctionStatus.clear"]();
+      mockStatusRecorder.clear();
     }
   }
 
diff --git a/src/core/executionRecorder.ts b/src/core/executionRecorder.ts
@@ -107,26 +107,24 @@ export class ExecutionRecorder implements UnitTestFramework {
     this.logRecorder.addLog(msg);
   }
 
-  getCollectionFuncSet(arg: ImportsArgument): Record<string, Record<string, unknown>> {
+  getCollectionFuncSet(arg: ImportsArgument): Record<string, unknown> {
     return {
-      __unittest_framework_env: {
-        addDescription: (description: number): void => {
-          this._addDescription(arg.exports!.__getString(description));
-        },
-        removeDescription: (): void => {
-          this._removeDescription();
-        },
-        registerTestFunction: (index: number): void => {
-          this._registerTestFunction(index);
-        },
-        collectCheckResult: (result: number, codeInfoIndex: number, actualValue: number, expectValue: number): void => {
-          this.collectCheckResult(
-            result !== 0,
-            codeInfoIndex,
-            arg.exports!.__getString(actualValue),
-            arg.exports!.__getString(expectValue)
-          );
-        },
+      addDescription: (description: number): void => {
+        this._addDescription(arg.exports!.__getString(description));
+      },
+      removeDescription: (): void => {
+        this._removeDescription();
+      },
+      registerTestFunction: (index: number): void => {
+        this._registerTestFunction(index);
+      },
+      collectCheckResult: (result: number, codeInfoIndex: number, actualValue: number, expectValue: number): void => {
+        this.collectCheckResult(
+          result !== 0,
+          codeInfoIndex,
+          arg.exports!.__getString(actualValue),
+          arg.exports!.__getString(expectValue)
+        );
       },
     };
   }
diff --git a/src/core/instrument.ts b/src/core/instrument.ts
@@ -16,7 +16,7 @@ export async function instrument(
     const baseName = sourceFile.slice(0, -5);
     const result = new InstrumentResult(baseName);
 
-    const reportFunction = "covInstrument/traceExpression";
+    const reportFunction = "__unittest_framework_env/traceExpression";
 
     const source = instrumenter.allocateUTF8(sourceFile);
     const output = instrumenter.allocateUTF8(result.instrumentedWasm);
diff --git a/src/core/mockStatusRecorder.ts b/src/core/mockStatusRecorder.ts
@@ -0,0 +1,74 @@
+interface MockObject {
+  calls: number;
+  ignore: boolean;
+  mockFunctionIndex: number;
+}
+
+export class MockStatusRecorder {
+  private _mockStatus = new Map<number, MockObject>();
+
+  private hasMocked(functionIndex: number): boolean {
+    const mockObject = this._mockStatus.get(functionIndex);
+    if (mockObject === undefined) {
+      return false;
+    }
+    return !mockObject.ignore;
+  }
+
+  // isCall = true,  return -1 if not mocked;
+  // isCall = false, return oldIndex if not mocked.
+  _checkMock(functionIndex: number, isCall: boolean): number {
+    if (this.hasMocked(functionIndex)) {
+      const mockObject = this._mockStatus.get(functionIndex);
+      mockObject!.calls++;
+      return mockObject!.mockFunctionIndex;
+    }
+    return isCall ? -1 : functionIndex;
+  }
+
+  _setMockFunction(originalFunctionIndex: number, mockFunctionIndex: number): void {
+    const mockObject: MockObject = {
+      calls: 0,
+      ignore: false,
+      mockFunctionIndex,
+    };
+    this._mockStatus.set(originalFunctionIndex, mockObject);
+  }
+
+  _getMockedFunctionCalls(originalFunctionIndex: number, mockFunctionIndex: number): number {
+    const mockObject = this._mockStatus.get(originalFunctionIndex);
+    if (mockObject === undefined || mockObject.mockFunctionIndex !== mockFunctionIndex) {
+      return 0;
+    }
+    return mockObject.calls;
+  }
+
+  _setMockedFunctionIgnore(originalFunctionIndex: number, ignore: boolean): void {
+    const mockObject = this._mockStatus.get(originalFunctionIndex);
+    if (mockObject === undefined) {
+      return;
+    }
+    mockObject.ignore = ignore;
+  }
+
+  clear(): void {
+    this._mockStatus.clear();
+  }
+
+  getMockFuncSet(): Record<string, unknown> {
+    return {
+      checkMock: (functionIndex: number, isCall: boolean): number => {
+        return this._checkMock(functionIndex, isCall);
+      },
+      setMockFunction: (originalFunctionIndex: number, mockFunctionIndex: number): void => {
+        this._setMockFunction(originalFunctionIndex, mockFunctionIndex);
+      },
+      getMockedFunctionCalls: (originalFunctionIndex: number, mockFunctionIndex: number): number => {
+        return this._getMockedFunctionCalls(originalFunctionIndex, mockFunctionIndex);
+      },
+      setMockedFunctionIgnore: (originalFunctionIndex: number, ignore: boolean): void => {
+        this._setMockedFunctionIgnore(originalFunctionIndex, ignore);
+      },
+    };
+  }
+}
diff --git a/src/utils/import.ts b/src/utils/import.ts
diff --git a/tests/cpp/lit/run.cjs b/tests/cpp/lit/run.cjs
diff --git a/tests/e2e/assertFailed/stdout.txt b/tests/e2e/assertFailed/stdout.txt
diff --git a/tests/ts/test/core/mockStatusRecorder.test.ts b/tests/ts/test/core/mockStatusRecorder.test.ts
diff --git a/tests/ts/test/utils/import.test.ts b/tests/ts/test/utils/import.test.ts

Original file line number	Diff line number	Diff line change
`@@ -26,7 +26,7 @@ void CovInstrumentationWalker::introduceReportFun() noexcept {`
`26`	`26`	`std::array<BinaryenType, 3U> iii_{BinaryenTypeInt32(), BinaryenTypeInt32(),`
`27`	`27`	`BinaryenTypeInt32()};`
`28`	`28`	`const BinaryenType iii = BinaryenTypeCreate(iii_.data(), iii_.size());`
`29`		`- BinaryenAddFunctionImport(module, reportFunName, "covInstrument", "traceExpression", iii,`
	`29`	`+ BinaryenAddFunctionImport(module, reportFunName, "__unittest_framework_env", "traceExpression", iii,`
`30`	`30`	`wasm::Type::none);`
`31`	`31`	`}`
`32`	`32`	`}`
Original file line number	Diff line number	Diff line change
`@@ -83,7 +83,7 @@ bool MockInstrumentationWalker::mockFunctionDuplicateImportedCheck() const noexc`
`83`	`83`	`std::array<BinaryenType, 2U> ii_ =`
`84`	`84`	`std::array<BinaryenType, 2U>{BinaryenTypeInt32(), BinaryenTypeInt32()};`
`85`	`85`	`const BinaryenType ii = BinaryenTypeCreate(ii_.data(), ii_.size());`
`86`		`- BinaryenAddFunctionImport(module, this->checkMock.data(), "mockInstrument", "checkMock", ii,`
	`86`	`+ BinaryenAddFunctionImport(module, this->checkMock.data(), "__unittest_framework_env", "checkMock", ii,`
`87`	`87`	`BinaryenTypeInt32());`
`88`	`88`	`}`
`89`	`89`	`return checkRepeat;`