feat: add call to internal state variables in mocks

defi-wonderland · Apr 30, 2024 · ba00d15 · ba00d15
1 parent ff8a692
commit ba00d15
Show file tree

Hide file tree

Showing 6 changed files with 229 additions and 18 deletions.
diff --git a/solidity/contracts/utils/ContractG.sol b/solidity/contracts/utils/ContractG.sol
@@ -37,7 +37,11 @@ contract ContractG {
     CommonStruct _common;
   }
 
-  mapping(bytes32 _disputeId => bool _finished) public _finished;
+  mapping(bytes32 _disputeId => bool _finished) public finished;
+
+  mapping(bytes32 _disputeId => bool _finished) internal _finishedInternal;
+
+  mapping(bytes32 _key1 => mapping(bytes32 _key2 => bool _finished)) internal _doubleFinishedInternal;
 
   mapping(bytes32 _disputeId => Set _votersSet) internal _votersA;
 
@@ -51,10 +55,14 @@ contract ContractG {
 
   mapping(bytes32 _disputeId => NestedStruct _nestedStruct) public _nestedStructs;
 
+  mapping(bytes32 _disputeId => NestedStruct _nestedStruct) internal _nestedStructsInternal;
+
   NestedStruct public nestedStruct;
 
   CommonStruct[] public structArray;
 
+  CommonStruct[] internal _structArrayInternal;
+
   CommonStruct[][] public twoDimensionalStruct;
 
   CommonStruct[][][] public threeDimensionalStruct;

diff --git a/solidity/test/utils/ContractD.t.sol b/solidity/test/utils/ContractD.t.sol
@@ -0,0 +1,63 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+import {Test} from 'forge-std/Test.sol';
+import {MockContractD} from 'test/smock/contracts/utils/MockContractD.sol';
+import {SmockHelper} from 'test/smock/SmockHelper.sol';
+
+contract UnitMockContractD is Test, SmockHelper {
+  address internal _owner = makeAddr('owner');
+  MockContractD internal _contractTest;
+
+  uint256 internal _initialValue = 5;
+  uint256 internal _newValue = 10;
+  string internal _someText = 'some text';
+
+  function setUp() public {
+    vm.prank(_owner);
+
+    _contractTest =
+      MockContractD(deployMock('TestContractD', type(MockContractD).creationCode, abi.encode(_initialValue)));
+  }
+
+  function test_Set_InternalUintVar() public {
+    _contractTest.set__internalUintVar(_newValue);
+    assertEq(_contractTest.call__internalUintVar(), _newValue);
+  }
+
+  function test_Call_InternalUintVar() public {
+    _contractTest.expectCall__setInternalUintVar(_initialValue);
+    _contractTest.mock_call__setInternalUintVar(_initialValue, true, _newValue, _someText);
+
+    (bool _success, uint256 _value, string memory _text) = _contractTest.call__setInternalUintVar(_initialValue);
+
+    assert(_success);
+    assertEq(_value, _newValue);
+    assertEq(_text, _someText);
+  }
+
+  function test_Call_GetVariables() public {
+    _contractTest.expectCall__getVariables(_initialValue);
+    _contractTest.mock_call__getVariables(_initialValue, true, _newValue, _someText);
+
+    (bool _success, uint256 _value, string memory _text) = _contractTest.call__getVariables(_initialValue);
+
+    assert(_success);
+    assertEq(_value, _newValue);
+    assertEq(_text, _someText);
+  }
+
+  function test_Call_InternalFuncNoInputNoOutput() public {
+    _contractTest.expectCall__internalFuncNoInputNoOutput();
+    _contractTest.mock_call__internalFuncNoInputNoOutput();
+
+    _contractTest.call__internalFuncNoInputNoOutput();
+  }
+
+  function test_Call_InternalViewFuncNoInputNoOutput() public {
+    _contractTest.expectCall__internalViewFuncNoInputNoOutput();
+    _contractTest.mock_call__internalViewFuncNoInputNoOutput();
+
+    _contractTest.call__internalViewFuncNoInputNoOutput();
+  }
+}
diff --git a/solidity/test/utils/ContractG.t.sol b/solidity/test/utils/ContractG.t.sol
@@ -227,4 +227,126 @@ contract UnitMockContractG is Test, SmockHelper {
     (_value) = _contractTest.twoDimensionalStringArray(1, 1);
     assertEq(_value, '40');
   }
+
+  function test_Set_NestedStruct() public {
+    ContractG.CommonStruct memory _commonStruct = ContractG.CommonStruct(20);
+    ContractG.NestedStruct memory _nestedStruct = ContractG.NestedStruct(10, _commonStruct);
+
+    _contractTest.set_nestedStruct(_nestedStruct);
+
+    (uint256 _counter, ContractG.CommonStruct memory _commonResult) = _contractTest.nestedStruct();
+
+    assertEq(_counter, 10);
+    assertEq(_commonResult._value, 20);
+  }
+
+  function test_Call_NestedStruct() public {
+    ContractG.CommonStruct memory _commonStruct = ContractG.CommonStruct(20);
+    ContractG.NestedStruct memory _nestedStruct = ContractG.NestedStruct(10, _commonStruct);
+
+    _contractTest.mock_call_nestedStruct(_nestedStruct);
+
+    (uint256 _counter, ContractG.CommonStruct memory _commonResult) = _contractTest.nestedStruct();
+
+    assertEq(_counter, 10);
+    assertEq(_commonResult._value, 20);
+  }
+
+  function test_Call_SetNestedStruct() public {
+    ContractG.CommonStruct memory _commonStruct = ContractG.CommonStruct(20);
+    ContractG.NestedStruct memory _nestedStruct = ContractG.NestedStruct(10, _commonStruct);
+
+    _contractTest.mock_call_setNestedStruct(_nestedStruct);
+    _contractTest.setNestedStruct(_nestedStruct);
+
+    (uint256 _counter, ContractG.CommonStruct memory _commonResult) = _contractTest.nestedStruct();
+
+    assertEq(_counter, 10);
+    assertEq(_commonResult._value, 20);
+  }
+
+  function test_Set_StructArray() public {
+    ContractG.CommonStruct[] memory _structArray = new ContractG.CommonStruct[](2);
+    _structArray[0] = ContractG.CommonStruct(10);
+    _structArray[1] = ContractG.CommonStruct(20);
+
+    _contractTest.set_structArray(_structArray);
+    (uint256 _value) = _contractTest.structArray(0);
+    assertEq(_value, 10);
+
+    (_value) = _contractTest.structArray(1);
+    assertEq(_value, 20);
+  }
+
+  function test_Call_StructArray() public {
+    ContractG.CommonStruct[] memory _structArray = new ContractG.CommonStruct[](2);
+    _structArray[0] = ContractG.CommonStruct(10);
+    _structArray[1] = ContractG.CommonStruct(20);
+
+    _contractTest.mock_call_structArray(0, _structArray[0]);
+    (uint256 _value) = _contractTest.structArray(0);
+    assertEq(_value, 10);
+
+    _contractTest.mock_call_structArray(1, _structArray[1]);
+    (_value) = _contractTest.structArray(1);
+    assertEq(_value, 20);
+  }
+
+  function test_Set_StructArrayInternal() public {
+    ContractG.CommonStruct[] memory _structArray = new ContractG.CommonStruct[](2);
+    _structArray[0] = ContractG.CommonStruct(10);
+    _structArray[1] = ContractG.CommonStruct(20);
+
+    _contractTest.set__structArrayInternal(_structArray);
+
+    ContractG.CommonStruct[] memory _resultArray = _contractTest.call__structArrayInternal();
+    assertEq(_resultArray[0]._value, 10);
+    assertEq(_resultArray[1]._value, 20);
+  }
+
+  function test_Set_Finished() public {
+    bytes32 _key = 'key';
+
+    _contractTest.set_finished(_key, true);
+
+    bool _value = _contractTest.finished(_key);
+
+    assertTrue(_value);
+  }
+
+  function test_Call_Finished() public {
+    bytes32 _key = 'key';
+    _contractTest.mock_call_finished(_key, true);
+
+    bool _value = _contractTest.finished(_key);
+    assertTrue(_value);
+  }
+
+  function test_Set_FinishedInternal() public {
+    bytes32 _key = 'key';
+    _contractTest.set__finishedInternal(_key, true);
+
+    bool _value = _contractTest.call__finishedInternal(_key);
+    assertTrue(_value);
+  }
+
+  function test_Set_DoubleFinishedInternal() public {
+    bytes32 _key0 = 'key0';
+    bytes32 _key1 = 'key1';
+    _contractTest.set__doubleFinishedInternal(_key0, _key1, true);
+
+    bool _value = _contractTest.call__doubleFinishedInternal(_key0, _key1);
+    assertTrue(_value);
+  }
+
+  function test_Set_NestedStructsInternal() public {
+    bytes32 _key = 'key';
+    ContractG.NestedStruct memory _nestedStruct = ContractG.NestedStruct(10, ContractG.CommonStruct(20));
+
+    _contractTest.set__nestedStructsInternal(_key, _nestedStruct);
+
+    ContractG.NestedStruct memory _result = _contractTest.call__nestedStructsInternal(_key);
+    assertEq(_result._counter, 10);
+    assertEq(_result._common._value, 20);
+  }
 }
diff --git a/src/templates/partials/array-state-variable.hbs b/src/templates/partials/array-state-variable.hbs
@@ -10,22 +10,28 @@
   }
 
   {{#unless isInternal}}
-  function mock_call_{{mockFunction.functionName}}({{#each dimensions}}uint256 _index{{@index}}, {{/each}} {{mockFunction.baseType}} _value) public {
-    vm.mockCall(
-      address(this),
-      abi.encodeWithSignature(
-        '{{mockFunction.functionName}}({{#each dimensions}}uint256{{#unless @last}},{{/unless}}{{/each}})', 
-        {{#each dimensions}}_index{{@index}}{{#unless @last}}, {{/unless}} {{/each}}),
-      abi.encode(
-        {{#if isStructArray}}
-          {{#each mockFunction.structFields}}
-            _value.{{this}}{{#unless @last}}, {{/unless}}
-          {{/each}}
-        {{else}}
-          _value
-        {{/if}}
-      )
-    );
-  }
+    function mock_call_{{mockFunction.functionName}}({{#each dimensions}}uint256 _index{{@index}}, {{/each}} {{mockFunction.baseType}} _value) public {
+      vm.mockCall(
+        address(this),
+        abi.encodeWithSignature(
+          '{{mockFunction.functionName}}({{#each dimensions}}uint256{{#unless @last}},{{/unless}}{{/each}})', 
+          {{#each dimensions}}_index{{@index}}{{#unless @last}}, {{/unless}} {{/each}}),
+        abi.encode(
+          {{#if isStructArray}}
+            {{#each mockFunction.structFields}}
+              _value.{{this}}{{#unless @last}}, {{/unless}}
+            {{/each}}
+          {{else}}
+            _value
+          {{/if}}
+        )
+      );
+    }
   {{/unless}}
+
+  {{#if isInternal}}
+    function call_{{setFunction.functionName}}() view public returns ({{setFunction.arrayType}}) {
+      return {{setFunction.functionName}};
+    }
+  {{/if}}
 {{/unless}}
diff --git a/src/templates/partials/mapping-state-variable.hbs b/src/templates/partials/mapping-state-variable.hbs
@@ -35,4 +35,10 @@
       );
     }
   {{/unless}}
+
+  {{#if isInternal}}
+    function call_{{setFunction.functionName}}({{#each setFunction.keyTypes}}{{this}} _key{{@index}}{{#unless @last}}, {{/unless}}{{/each}}) view public returns ({{setFunction.valueType}}) {
+      return {{setFunction.functionName}}{{#each setFunction.keyTypes}}[_key{{@index}}]{{/each}};
+    }
+  {{/if}}
 {{/unless}}
diff --git a/src/templates/partials/state-variable.hbs b/src/templates/partials/state-variable.hbs
@@ -19,3 +19,9 @@ function set_{{setFunction.functionName}}({{setFunction.paramType}} _{{setFuncti
     );
   }
 {{/unless}}
+
+{{#if isInternal}}
+  function call_{{setFunction.functionName}}() view public returns ({{setFunction.paramType}}) {
+    return {{setFunction.functionName}};
+  }
+{{/if}}