test(contracts): fix invariant tests with retreive eth

nocturne-xyz · Aug 7, 2023 · 62da0dd · 62da0dd
1 parent 1731059
commit 62da0dd
Show file tree

Hide file tree

Showing 4 changed files with 75 additions and 16 deletions.
diff --git a/packages/contracts/contracts/DepositManager.sol b/packages/contracts/contracts/DepositManager.sol
@@ -354,13 +354,13 @@ contract DepositManager is
         _weth.deposit{value: totalDepositAmount}();
     }
 
-    /// @notice Retrieves an ETH deposit either prematurely because user cancelled or because 
+    /// @notice Retrieves an ETH deposit either prematurely because user cancelled or because
     ///         screener didn't complete. Unwraps weth back to eth and send to user.
     /// @dev We accept race condition where user could technically retrieve their deposit before
     ///      the screener completes it. This would grief the screener but would incur a greater
     ///      cost to the user to continually instantiate + prematurely retrieve.
-    /// @dev Same code as normal retrieveDeposit except deposit is checked to be weth and weth is 
-    //       unwrapped to eth. The call to send gas compensation eth back to user also sends back 
+    /// @dev Same code as normal retrieveDeposit except deposit is checked to be weth and weth is
+    //       unwrapped to eth. The call to send gas compensation eth back to user also sends back
     ///      deposit.value eth.
     /// @param req Deposit request corresponding to ETH deposit to retrieve
     function retrieveETHDeposit(

diff --git a/packages/contracts/contracts/test/invariant/InvariantsBase.sol b/packages/contracts/contracts/test/invariant/InvariantsBase.sol
@@ -119,20 +119,25 @@ contract InvariantsBase is Test {
     function assert_deposit_depositManagerBalanceEqualsInMinusOutErc20s()
         internal
     {
+        uint256 totalETHRetrieved = depositManagerHandler
+            .ghost_retrieveDepositSumETH();
+
         for (uint256 i = 0; i < depositErc20s.length; i++) {
+            uint256 expectedErc20Balance = depositManagerHandler
+                .ghost_instantiateDepositSumErc20ForToken(i) -
+                depositManagerHandler.ghost_retrieveDepositSumErc20ForToken(i) -
+                depositManagerHandler.ghost_completeDepositSumErc20ForToken(i);
+
+            // WETH case
+            if (i == 0) {
+                expectedErc20Balance -= totalETHRetrieved;
+            }
+
             assertEq(
                 IERC20(depositErc20s[i]).balanceOf(
                     address(depositManagerHandler.depositManager())
                 ),
-                depositManagerHandler.ghost_instantiateDepositSumErc20ForToken(
-                    i
-                ) -
-                    depositManagerHandler.ghost_retrieveDepositSumErc20ForToken(
-                        i
-                    ) -
-                    depositManagerHandler.ghost_completeDepositSumErc20ForToken(
-                        i
-                    )
+                expectedErc20Balance
             );
         }
     }
@@ -207,10 +212,16 @@ contract InvariantsBase is Test {
         for (uint256 i = 0; i < allActors.length; i++) {
             uint256 actorBalance = allActors[i].balance;
             console.logUint(actorBalance);
-            uint256 actorBalanceCap = depositManagerHandler
+            uint256 actorBalanceETHRetrieved = depositManagerHandler
+                .ghost_retrieveDepositSumETHForActor(allActors[i]);
+            uint256 actorMaxGasCompRetrieved = depositManagerHandler
                 .ghost_totalSuppliedGasCompensationForActor(allActors[i]);
 
-            assertLe(actorBalance, actorBalanceCap);
+            assertGe(actorBalance, actorBalanceETHRetrieved);
+            assertLe(
+                actorBalance,
+                actorBalanceETHRetrieved + actorMaxGasCompRetrieved
+            );
         }
     }
 

diff --git a/packages/contracts/contracts/test/invariant/ProtocolInvariants.t.sol b/packages/contracts/contracts/test/invariant/ProtocolInvariants.t.sol
@@ -147,17 +147,20 @@ contract ProtocolInvariants is Test, InvariantsBase {
             depositErc20
         );
 
-        bytes4[] memory depositManagerHandlerSelectors = new bytes4[](4);
+        bytes4[] memory depositManagerHandlerSelectors = new bytes4[](5);
         depositManagerHandlerSelectors[0] = depositManagerHandler
             .instantiateDepositETH
             .selector;
         depositManagerHandlerSelectors[1] = depositManagerHandler
             .instantiateDepositErc20
             .selector;
         depositManagerHandlerSelectors[2] = depositManagerHandler
-            .retrieveDepositErc20
+            .retrieveDepositETH
             .selector;
         depositManagerHandlerSelectors[3] = depositManagerHandler
+            .retrieveDepositErc20
+            .selector;
+        depositManagerHandlerSelectors[4] = depositManagerHandler
             .completeDepositErc20
             .selector;
 

diff --git a/packages/contracts/contracts/test/invariant/actors/DepositManagerHandler.sol b/packages/contracts/contracts/test/invariant/actors/DepositManagerHandler.sol
@@ -54,6 +54,7 @@ contract DepositManagerHandler is CommonBase, StdCheats, StdUtils {
 
     // First entry in array is for weth
     ActorSumSet[] internal _instantiateDepositSumSetErc20s;
+    ActorSumSet internal _retrieveDepositSumSetETH;
     ActorSumSet[] internal _retrieveDepositSumSetErc20s;
     ActorSumSet[] internal _completeDepositSumSetErc20s;
 
@@ -143,6 +144,7 @@ contract DepositManagerHandler is CommonBase, StdCheats, StdUtils {
             "retrieveDepositSumErc20",
             ghost_retrieveDepositSumErc20ForToken(1)
         );
+        console.log("retrieveDepositSumETH", ghost_retrieveDepositSumETH());
         console.log(
             "completeDepositSumErc20",
             ghost_completeDepositSumErc20ForToken(1)
@@ -311,6 +313,39 @@ contract DepositManagerHandler is CommonBase, StdCheats, StdUtils {
         );
     }
 
+    function retrieveDepositETH(
+        uint256 seed
+    ) public trackCall("retrieveDepositETH") {
+        // Get random request
+        uint256 index;
+        if (_depositSet.length > 0) {
+            index = seed % _depositSet.length;
+        } else {
+            return;
+        }
+
+        DepositRequest memory randDepositRequest = _depositSet[index];
+
+        (, address assetAddr, ) = AssetUtils.decodeAsset(
+            randDepositRequest.encodedAsset
+        );
+        if (assetAddr != address(depositManager._weth())) {
+            lastCall = "no-op";
+            return;
+        }
+
+        // Retrieve deposit
+        vm.prank(randDepositRequest.spender);
+        try depositManager.retrieveETHDeposit(randDepositRequest) {
+            _retrieveDepositSumSetETH.addToActorSum(
+                randDepositRequest.spender,
+                randDepositRequest.value
+            );
+        } catch {
+            _reverts["retrieveDepositETH"] += 1;
+        }
+    }
+
     function retrieveDepositErc20(
         uint256 seed
     ) public trackCall("retrieveDepositErc20") {
@@ -440,6 +475,10 @@ contract DepositManagerHandler is CommonBase, StdCheats, StdUtils {
         return _retrieveDepositSumSetErc20s[tokenIndex].getTotalForAll();
     }
 
+    function ghost_retrieveDepositSumETH() public view returns (uint256) {
+        return _retrieveDepositSumSetETH.getTotalForAll();
+    }
+
     function ghost_completeDepositSumErc20ForToken(
         uint256 tokenIndex
     ) public view returns (uint256) {
@@ -454,6 +493,12 @@ contract DepositManagerHandler is CommonBase, StdCheats, StdUtils {
             _instantiateDepositSumSetErc20s[tokenIndex].getSumForActor(actor);
     }
 
+    function ghost_retrieveDepositSumETHForActor(
+        address actor
+    ) public view returns (uint256) {
+        return _retrieveDepositSumSetETH.getSumForActor(actor);
+    }
+
     function ghost_retrieveDepositSumErc20ForActorOfToken(
         address actor,
         uint256 tokenIndex