Skip to content

Commit

Permalink
fix: mock internal functions without parameters (#33)
Browse files Browse the repository at this point in the history
  • Loading branch information
0xJabberwock authored Dec 15, 2023
1 parent 431a158 commit ddb009d
Show file tree
Hide file tree
Showing 4 changed files with 40 additions and 29 deletions.
4 changes: 0 additions & 4 deletions solidity/contracts/utils/ContractB.sol
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,4 @@ contract ContractB is IContractB {
stringVariable = _stringVariable;
_result = true;
}

function _internalNoReturn() internal {
stringVariable = 'internalNoReturn';
}
}
8 changes: 8 additions & 0 deletions solidity/contracts/utils/ContractD.sol
Original file line number Diff line number Diff line change
Expand Up @@ -15,4 +15,12 @@ contract ContractD {
function _getVariables(uint256 _uintVariable) internal view virtual returns (bool, uint256, string memory) {
return (true, 111, 'test');
}

function _internalNoInputNoOutput() internal virtual {
_internalUintVar = 11;
}

function _internalViewNoInputNoOutput() internal view virtual {
uint256 __internalUintVar = _internalUintVar;
}
}
13 changes: 10 additions & 3 deletions src/templates/mockInternalFunctionTemplate.hbs
Original file line number Diff line number Diff line change
@@ -1,18 +1,22 @@
{{#if isView}}
{{#if outputs}}
struct {{functionName}}Output {
{{#each outputTypes}}
{{this}} {{lookup ../outputNames @index}};
{{/each}}
}

mapping(bytes32 => {{functionName}}Output) private {{functionName}}Outputs;
{{/if}}
bytes32[] private {{functionName}}InputHashes;
{{/if}}

function mock_call_{{functionName}}({{parameters}}) public {
{{#if isView}}
bytes32 _key = keccak256(abi.encode({{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}}));
{{#if outputs}}
{{functionName}}Outputs[_key] = {{functionName}}Output({{#each outputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}});
{{/if}}
for (uint256 _i; _i < {{functionName}}InputHashes.length; ++_i) {
if (_key == {{functionName}}InputHashes[_i]) {
return;
Expand All @@ -22,7 +26,7 @@ function mock_call_{{functionName}}({{parameters}}) public {
{{else}}
vm.mockCall(
address(this),
abi.encodeWithSignature("{{signature}}", {{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}}),
abi.encodeWithSignature("{{signature}}"{{#if inputs}}, {{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}}{{/if}}),
abi.encode({{#each outputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}})
);
{{/if}}
Expand All @@ -33,14 +37,17 @@ function {{functionName}}({{inputs}}) internal {{#if isView}}view {{/if}}overrid
bytes32 _key = keccak256(abi.encode({{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}}));
for (uint256 _i; _i < {{functionName}}InputHashes.length; ++_i) {
if (_key == {{functionName}}InputHashes[_i]) {
{{#if outputs}}
{{functionName}}Output memory _output = {{functionName}}Outputs[_key];
{{/if}}
return ({{#each outputNames}}_output.{{this}}{{#unless @last}}, {{/unless}}{{/each}});
}
}
return super.{{functionName}}({{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}});
{{else}}
(bool _success, bytes memory _data) = address(this).call(abi.encodeWithSignature("{{signature}}", {{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}}));
({{#each outputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}}) = _success ? abi.decode(_data, ({{#each outputTypes}}{{this}}{{#unless @last}}, {{/unless}}{{/each}})) : super.{{functionName}}({{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}});
(bool _success, bytes memory _data) = address(this).call(abi.encodeWithSignature("{{signature}}"{{#if inputs}}, {{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}}{{/if}}));
if (_success) return abi.decode(_data, ({{#each outputTypes}}{{this}}{{#unless @last}}, {{/unless}}{{/each}}));
else return super.{{functionName}}({{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}});
{{/if}}
}
{{nl}}
44 changes: 22 additions & 22 deletions test/e2e/get-internal-functions.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -97,18 +97,18 @@ describe('E2E: getInternalMockFunctions', () => {
expect(param1).to.not.be.undefined;
expect(param1?.typeDescriptions.typeString).to.equal('uint256');

const return0 = func.returnParameters.parameters.find((param) => param.name === '_returnParam0');
expect(return0).to.not.be.undefined;
expect(return0?.typeDescriptions.typeString).to.equal('bool');

const return1 = func.returnParameters.parameters.find((param) => param.name === '_returnParam1');
expect(return1).to.not.be.undefined;
expect(return1?.typeDescriptions.typeString).to.equal('uint256');

const return2 = func.returnParameters.parameters.find((param) => param.name === '_returnParam2');
expect(return2).to.not.be.undefined;
expect(return2?.typeDescriptions.typeString).to.equal('string');
expect(return2?.storageLocation).to.equal('memory');
const returnParam1 = func.returnParameters.parameters.find((param) => param.name === '_returnParam0');
expect(returnParam1).to.not.be.undefined;
expect(returnParam1?.typeDescriptions.typeString).to.equal('bool');

const returnParam2 = func.returnParameters.parameters.find((param) => param.name === '_returnParam1');
expect(returnParam2).to.not.be.undefined;
expect(returnParam2?.typeDescriptions.typeString).to.equal('uint256');

const returnParam3 = func.returnParameters.parameters.find((param) => param.name === '_returnParam2');
expect(returnParam3).to.not.be.undefined;
expect(returnParam3?.typeDescriptions.typeString).to.equal('string');
expect(returnParam3?.storageLocation).to.equal('memory');
});

it('MockContractD must include internal view function mock call', async () => {
Expand Down Expand Up @@ -150,17 +150,17 @@ describe('E2E: getInternalMockFunctions', () => {
expect(param1).to.not.be.undefined;
expect(param1?.typeDescriptions.typeString).to.equal('uint256');

const return0 = func.returnParameters.parameters.find((param) => param.name === '_returnParam0');
expect(return0).to.not.be.undefined;
expect(return0?.typeDescriptions.typeString).to.equal('bool');
const returnParam1 = func.returnParameters.parameters.find((param) => param.name === '_returnParam0');
expect(returnParam1).to.not.be.undefined;
expect(returnParam1?.typeDescriptions.typeString).to.equal('bool');

const return1 = func.returnParameters.parameters.find((param) => param.name === '_returnParam1');
expect(return1).to.not.be.undefined;
expect(return1?.typeDescriptions.typeString).to.equal('uint256');
const returnParam2 = func.returnParameters.parameters.find((param) => param.name === '_returnParam1');
expect(returnParam2).to.not.be.undefined;
expect(returnParam2?.typeDescriptions.typeString).to.equal('uint256');

const return2 = func.returnParameters.parameters.find((param) => param.name === '_returnParam2');
expect(return2).to.not.be.undefined;
expect(return2?.typeDescriptions.typeString).to.equal('string');
expect(return2?.storageLocation).to.equal('memory');
const returnParam3 = func.returnParameters.parameters.find((param) => param.name === '_returnParam2');
expect(returnParam3).to.not.be.undefined;
expect(returnParam3?.typeDescriptions.typeString).to.equal('string');
expect(returnParam3?.storageLocation).to.equal('memory');
});
});

0 comments on commit ddb009d

Please sign in to comment.