feat: add call to internal functions in mock (#70)

defi-wonderland · Apr 24, 2024 · 9dbf910 · 9dbf910
1 parent ebb0895
commit 9dbf910
Show file tree

Hide file tree

Showing 6 changed files with 56 additions and 8 deletions.
diff --git a/solidity/contracts/utils/ContractG.sol b/solidity/contracts/utils/ContractG.sol
@@ -52,4 +52,10 @@ contract ContractG {
   mapping(bytes32 _disputeId => NestedStruct _nestedStruct) public _nestedStructs;
 
   NestedStruct public nestedStruct;
+
+  // CommonStruct[][] public multidimensionalStruct;
+
+  function setNestedStruct(NestedStruct memory _nestedStruct) public {
+    nestedStruct = _nestedStruct;
+  }
 }
diff --git a/solidity/contracts/utils/ContractJ.sol b/solidity/contracts/utils/ContractJ.sol
@@ -28,10 +28,4 @@ contract ContractJ {
     myUserTypeVariable = add(a, b);
     return myUserTypeVariable;
   }
-
-  // =============== virtual call ===============
-
-  function callInternalAdd(MyUserType a, MyUserType b) public returns (MyUserType) {
-    return internalAdd(a, b);
-  }
 }
diff --git a/solidity/test/ContractTest.t.sol b/solidity/test/ContractTest.t.sol
@@ -340,6 +340,19 @@ contract E2EMockContractTest_Mock_call_Internal_Func is CommonE2EBase {
     assertEq(_res3, 'test');
   }
 
+  function test_call_InternalVirtualFunction() public {
+    _contractTest.mock_call_internalVirtualFunction(10, false, 12, 'TEST');
+    (bool _res1, uint256 _res2, string memory _res3) = _contractTest.call_internalVirtualFunction(10);
+    assertEq(_res1, false);
+    assertEq(_res2, 12);
+    assertEq(_res3, 'TEST');
+
+    (_res1, _res2, _res3) = _contractTest.call_internalVirtualFunction(11);
+    assertEq(_res1, true);
+    assertEq(_res2, 1);
+    assertEq(_res3, 'test');
+  }
+
   function test_MockCall_InternalViewVirtualFunction() public {
     _contractTest.mock_call_internalViewVirtualFunction(10, false, 12, 'TEST');
     (bool _res1, uint256 _res2, string memory _res3) = _contractTest.callInternalViewVirtualFunction(10);
@@ -353,6 +366,19 @@ contract E2EMockContractTest_Mock_call_Internal_Func is CommonE2EBase {
     assertEq(_res3, 'test');
   }
 
+  function test_call_InternalViewVirtualFunction() public {
+    _contractTest.mock_call_internalViewVirtualFunction(10, false, 12, 'TEST');
+    (bool _res1, uint256 _res2, string memory _res3) = _contractTest.call_internalViewVirtualFunction(10);
+    assertEq(_res1, false);
+    assertEq(_res2, 12);
+    assertEq(_res3, 'TEST');
+
+    (_res1, _res2, _res3) = _contractTest.call_internalViewVirtualFunction(11);
+    assertEq(_res1, true);
+    assertEq(_res2, 1);
+    assertEq(_res3, 'test');
+  }
+
   function test_MockCall_InternalPureVirtualFunction() public {
     _contractTest.mock_call_internalPureVirtualFunction(10, false, 12, 'TEST');
     (bool _res1, uint256 _res2, string memory _res3) = _contractTest.callInternalPureVirtualFunction(10);
@@ -365,4 +391,17 @@ contract E2EMockContractTest_Mock_call_Internal_Func is CommonE2EBase {
     assertEq(_res2, 11);
     assertEq(_res3, 'test');
   }
+
+  function test_call_InternalPureVirtualFunction() public {
+    _contractTest.mock_call_internalPureVirtualFunction(10, false, 12, 'TEST');
+    (bool _res1, uint256 _res2, string memory _res3) = _contractTest.call_internalPureVirtualFunction(10);
+    assertEq(_res1, false);
+    assertEq(_res2, 12);
+    assertEq(_res3, 'TEST');
+
+    (_res1, _res2, _res3) = _contractTest.call_internalPureVirtualFunction(11);
+    assertEq(_res1, true);
+    assertEq(_res2, 11);
+    assertEq(_res3, 'test');
+  }
 }
diff --git a/solidity/test/utils/ContractJ.t.sol b/solidity/test/utils/ContractJ.t.sol
@@ -88,7 +88,7 @@ contract UnitMockContractJ is Test, SmockHelper {
 
   function test_InternalAdd() public {
     _contractTest.mock_call_internalAdd(_myUserTypeA, _myUserTypeB, _myUserTypeResult);
-    assertEq(ContractJ.MyUserType.unwrap(_contractTest.callInternalAdd(_myUserTypeA, _myUserTypeB)), _result);
+    assertEq(ContractJ.MyUserType.unwrap(_contractTest.call_internalAdd(_myUserTypeA, _myUserTypeB)), _result);
   }
 
   function test_ExternalAdd() public {

diff --git a/src/templates/partials/internal-function.hbs b/src/templates/partials/internal-function.hbs
@@ -15,3 +15,8 @@ function {{functionName}}({{inputs}}) internal override {{#if outputs}}returns (
     else return super.{{functionName}}({{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}});
   {{/if}}
 }
+
+
+function call_{{functionName}}({{inputs}}) public {{#if outputs}}returns ({{outputs}}){{/if}} {
+  return {{functionName}}({{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}});
+}
diff --git a/src/templates/partials/internal-view-function.hbs b/src/templates/partials/internal-view-function.hbs
@@ -57,4 +57,8 @@ function {{functionName}}{{#if isPure}}Helper{{/if}}({{inputs}}) internal view {
     return _{{functionName}}CastToPure({{functionName}}Helper)({{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}});
   }
 
-{{/if}}
+{{/if}}
+
+function call_{{functionName}}({{inputs}}) public {{#if isView}}view{{else}}pure{{/if}} {{#if outputs}}returns ({{outputs}}){{/if}} {
+  return {{functionName}}({{#each inputNames}}{{this}}{{#unless @last}}, {{/unless}}{{/each}});
+}