Skip to content

Commit

Permalink
feat: generate mock helper
Browse files Browse the repository at this point in the history
  • Loading branch information
dristpunk committed Nov 2, 2023
1 parent 9901d0b commit accbdfd
Show file tree
Hide file tree
Showing 4 changed files with 49 additions and 10 deletions.
16 changes: 7 additions & 9 deletions solidity/test/ContractTest.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -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');
Expand All @@ -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')))
);
}
}

Expand Down
14 changes: 13 additions & 1 deletion src/mock-contract-generator.ts
Original file line number Diff line number Diff line change
@@ -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';
Expand Down Expand Up @@ -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);
Expand Down
20 changes: 20 additions & 0 deletions src/templates/mockHelperTemplate.hbs
Original file line number Diff line number Diff line change
@@ -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);
}
}
9 changes: 9 additions & 0 deletions src/utils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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');
Expand All @@ -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
Expand Down

0 comments on commit accbdfd

Please sign in to comment.