From accbdfdc23923669a0ff984e316dae992a228029 Mon Sep 17 00:00:00 2001
From: Dristpunk <dristpunk@gmail.com>
Date: Thu, 2 Nov 2023 15:38:13 +0300
Subject: [PATCH] feat: generate mock helper

---
 solidity/test/ContractTest.t.sol     | 16 +++++++---------
 src/mock-contract-generator.ts       | 14 +++++++++++++-
 src/templates/mockHelperTemplate.hbs | 20 ++++++++++++++++++++
 src/utils.ts                         |  9 +++++++++
 4 files changed, 49 insertions(+), 10 deletions(-)
 create mode 100644 src/templates/mockHelperTemplate.hbs

diff --git a/solidity/test/ContractTest.t.sol b/solidity/test/ContractTest.t.sol
index b83a5d1..d0c2580 100644
--- a/solidity/test/ContractTest.t.sol
+++ b/solidity/test/ContractTest.t.sol
@@ -5,8 +5,9 @@ import {Test} from 'forge-std/Test.sol';
 import {IERC20} from 'isolmate/interfaces/tokens/IERC20.sol';
 import {MockContractTest} from 'test/mock-contracts/contracts/MockContractTest.sol';
 import {console} from 'forge-std/console.sol';
+import {MockHelper} from 'test/mock-contracts/MockHelper.sol';
 
-contract CommonE2EBase is Test {
+contract CommonE2EBase is Test, MockHelper {
   uint256 internal constant _FORK_BLOCK = 15_452_788;
 
   address internal _user = makeAddr('user');
@@ -16,14 +17,11 @@ contract CommonE2EBase is Test {
   function setUp() public {
     vm.createSelectFork(vm.rpcUrl('mainnet'), _FORK_BLOCK);
     vm.prank(_owner);
-    _contractTest = new MockContractTest(
-            1,
-            "2",
-            true,
-            _owner,
-            bytes32("4")
-        );
-    vm.allowCheatcodes(address(_contractTest));
+    // abi.encode(1, '2', true, _owner, bytes32('4'))
+
+    _contractTest = MockContractTest(
+      deployMock('Test', type(MockContractTest).creationCode, abi.encode(1, '2', true, _owner, bytes32('4')))
+    );
   }
 }
 
diff --git a/src/mock-contract-generator.ts b/src/mock-contract-generator.ts
index f37f8a6..0174227 100644
--- a/src/mock-contract-generator.ts
+++ b/src/mock-contract-generator.ts
@@ -1,5 +1,11 @@
 import { getExternalMockFunctions, getInternalMockFunctions, getConstructor, getImports, getStateVariables, Ast } from './index';
-import { getSubDirNameFromPath, registerHandlebarsTemplates, getContractNamesAndFolders, compileSolidityFilesFoundry } from './utils';
+import {
+  getSubDirNameFromPath,
+  registerHandlebarsTemplates,
+  getContractNamesAndFolders,
+  compileSolidityFilesFoundry,
+  registerMockHelperTemplate,
+} from './utils';
 import Handlebars from 'handlebars';
 import { writeFileSync, existsSync, readdirSync } from 'fs';
 import { ensureDir, emptyDir } from 'fs-extra';
@@ -129,6 +135,12 @@ export const generateMockContracts = async (
       writeFileSync(`${contractFolder}/Mock${contractName}.sol`, cleanedCode);
     });
 
+    // Generate MockHelper contract
+    const mockHelperTemplateContent: string = registerMockHelperTemplate();
+    const mockHelperTemplate = Handlebars.compile(mockHelperTemplateContent);
+    const mockHelperCode: string = mockHelperTemplate({});
+    writeFileSync(`${generatedContractsDir}/MockHelper.sol`, mockHelperCode);
+
     console.log('Mock contracts generated successfully');
     // Compile the mock contracts
     await compileSolidityFilesFoundry(generatedContractsDir);
diff --git a/src/templates/mockHelperTemplate.hbs b/src/templates/mockHelperTemplate.hbs
new file mode 100644
index 0000000..ae565d5
--- /dev/null
+++ b/src/templates/mockHelperTemplate.hbs
@@ -0,0 +1,20 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+import {Test} from 'forge-std/Test.sol';
+
+contract MockHelper is Test {
+  function deployMock(
+    string memory _label,
+    bytes memory _creationCode,
+    bytes memory _encodedArgs
+  ) internal returns (address _deployed) {
+    bytes memory _bytecode = abi.encodePacked(_creationCode, _encodedArgs);
+    assembly {
+      mstore(0x0, _creationCode)
+      _deployed := create2(0, add(_bytecode, 0x20), mload(_bytecode), 'Wonderland')
+    }
+    vm.label(_deployed, _label);
+    vm.allowCheatcodes(_deployed);
+  }
+}
diff --git a/src/utils.ts b/src/utils.ts
index 6c2d021..4ed1325 100644
--- a/src/utils.ts
+++ b/src/utils.ts
@@ -66,6 +66,7 @@ export const registerHandlebarsTemplates = (): string => {
   const basicStateVariablesTemplatePath = resolve(__dirname, 'templates', 'mockBasicStateVariableTemplate.hbs');
   const arrayStateVariablesTemplatePath = resolve(__dirname, 'templates', 'mockArrayStateVariableTemplate.hbs');
   const mappingStateVariablesTemplatePath = resolve(__dirname, 'templates', 'mockMappingStateVariableTemplate.hbs');
+
   // Read the templates
   const templateContent = readFileSync(templatePath, 'utf8');
   const externalFuncsTemplateContent = readFileSync(externalFuncsTemplatePath, 'utf8');
@@ -80,9 +81,17 @@ export const registerHandlebarsTemplates = (): string => {
   Handlebars.registerPartial('mockInternalFunction', internalFuncsTemplateContent);
   Handlebars.registerPartial('mockArrayStateVariable', arrayStateVariablesTemplateContent);
   Handlebars.registerPartial('mockMappingStateVariable', mappingStateVariablesTemplateContent);
+
   return templateContent;
 };
 
+export const registerMockHelperTemplate = (): string => {
+  const mockHelperTemplatePath = resolve(__dirname, 'templates', 'mockHelperTemplate.hbs');
+  const mockHelperTemplateContent = readFileSync(mockHelperTemplatePath, 'utf8');
+  Handlebars.registerPartial('mockHelperTemplatet', mockHelperTemplateContent);
+  return mockHelperTemplateContent;
+};
+
 /**
  * Returns the names of the contracts in the given directory and its subdirectories
  * @param contractsDir The directory where the contracts are located