Add wrap up to functionality for WETH wrapper script (#104)

compound-finance · Nov 20, 2024 · 5b863e8 · 5b863e8
1 parent 30bb173
commit 5b863e8
Show file tree

Hide file tree

Showing 2 changed files with 114 additions and 0 deletions.
diff --git a/src/WrapperScripts.sol b/src/WrapperScripts.sol
@@ -10,10 +10,24 @@ contract WrapperActions {
         IWETH(weth).deposit{value: amount}();
     }
 
+    function wrapETHUpTo(address weth, uint256 targetAmount) external payable {
+        uint256 currentBalance = IERC20(weth).balanceOf(address(this));
+        if (currentBalance < targetAmount) {
+            IWETH(weth).deposit{value: targetAmount - currentBalance}();
+        }
+    }
+
     function unwrapWETH(address weth, uint256 amount) external {
         IWETH(weth).withdraw(amount);
     }
 
+    function unwrapWETHUpTo(address weth, uint256 targetAmount) external {
+        uint256 currentBalance = address(this).balance;
+        if (currentBalance < targetAmount) {
+            IWETH(weth).withdraw(targetAmount - currentBalance);
+        }
+    }
+
     function wrapLidoStETH(address wstETH, address stETH, uint256 amount) external {
         IERC20(stETH).approve(wstETH, amount);
         IWstETH(wstETH).wrap(amount);

diff --git a/test/WrapperScripts.t.sol b/test/WrapperScripts.t.sol
@@ -65,6 +65,56 @@ contract WrapperScriptsTest is Test {
         assertEq(address(wallet).balance, 0 ether);
     }
 
+    function testWrapETHUpTo() public {
+        vm.pauseGasMetering();
+        QuarkWallet wallet = QuarkWallet(factory.create(alice, address(0)));
+
+        deal(address(wallet), 10 ether);
+        deal(WETH, address(wallet), 7 ether);
+
+        QuarkWallet.QuarkOperation memory op = new QuarkOperationHelper().newBasicOpWithCalldata(
+            wallet,
+            wrapperScript,
+            abi.encodeWithSelector(WrapperActions.wrapETHUpTo.selector, WETH, 10 ether),
+            ScriptType.ScriptSource
+        );
+        bytes memory signature = new SignatureHelper().signOp(alicePrivateKey, wallet, op);
+
+        assertEq(IERC20(WETH).balanceOf(address(wallet)), 7 ether);
+        assertEq(address(wallet).balance, 10 ether);
+
+        vm.resumeGasMetering();
+        wallet.executeQuarkOperation(op, signature);
+
+        assertEq(IERC20(WETH).balanceOf(address(wallet)), 10 ether);
+        assertEq(address(wallet).balance, 7 ether);
+    }
+
+    function testWrapETHUpToDoesNotWrapIfNotNeeded() public {
+        vm.pauseGasMetering();
+        QuarkWallet wallet = QuarkWallet(factory.create(alice, address(0)));
+
+        deal(address(wallet), 10 ether);
+        deal(WETH, address(wallet), 10 ether);
+
+        QuarkWallet.QuarkOperation memory op = new QuarkOperationHelper().newBasicOpWithCalldata(
+            wallet,
+            wrapperScript,
+            abi.encodeWithSelector(WrapperActions.wrapETHUpTo.selector, WETH, 10 ether),
+            ScriptType.ScriptSource
+        );
+        bytes memory signature = new SignatureHelper().signOp(alicePrivateKey, wallet, op);
+
+        assertEq(IERC20(WETH).balanceOf(address(wallet)), 10 ether);
+        assertEq(address(wallet).balance, 10 ether);
+
+        vm.resumeGasMetering();
+        wallet.executeQuarkOperation(op, signature);
+
+        assertEq(IERC20(WETH).balanceOf(address(wallet)), 10 ether);
+        assertEq(address(wallet).balance, 10 ether);
+    }
+
     function testUnwrapWETH() public {
         vm.pauseGasMetering();
         QuarkWallet wallet = QuarkWallet(factory.create(alice, address(0)));
@@ -87,6 +137,56 @@ contract WrapperScriptsTest is Test {
         assertEq(address(wallet).balance, 10 ether);
     }
 
+    function testUnwrapWETHUpTo() public {
+        vm.pauseGasMetering();
+        QuarkWallet wallet = QuarkWallet(factory.create(alice, address(0)));
+
+        deal(WETH, address(wallet), 10 ether);
+        deal(address(wallet), 7 ether);
+
+        QuarkWallet.QuarkOperation memory op = new QuarkOperationHelper().newBasicOpWithCalldata(
+            wallet,
+            wrapperScript,
+            abi.encodeWithSelector(WrapperActions.unwrapWETHUpTo.selector, WETH, 10 ether),
+            ScriptType.ScriptSource
+        );
+        bytes memory signature = new SignatureHelper().signOp(alicePrivateKey, wallet, op);
+
+        assertEq(IERC20(WETH).balanceOf(address(wallet)), 10 ether);
+        assertEq(address(wallet).balance, 7 ether);
+
+        vm.resumeGasMetering();
+        wallet.executeQuarkOperation(op, signature);
+
+        assertEq(IERC20(WETH).balanceOf(address(wallet)), 7 ether);
+        assertEq(address(wallet).balance, 10 ether);
+    }
+
+    function testUnwrapWETHUpToDoesNotUnwrapIfNotNeeded() public {
+        vm.pauseGasMetering();
+        QuarkWallet wallet = QuarkWallet(factory.create(alice, address(0)));
+
+        deal(WETH, address(wallet), 10 ether);
+        deal(address(wallet), 10 ether);
+
+        QuarkWallet.QuarkOperation memory op = new QuarkOperationHelper().newBasicOpWithCalldata(
+            wallet,
+            wrapperScript,
+            abi.encodeWithSelector(WrapperActions.unwrapWETHUpTo.selector, WETH, 10 ether),
+            ScriptType.ScriptSource
+        );
+        bytes memory signature = new SignatureHelper().signOp(alicePrivateKey, wallet, op);
+
+        assertEq(IERC20(WETH).balanceOf(address(wallet)), 10 ether);
+        assertEq(address(wallet).balance, 10 ether);
+
+        vm.resumeGasMetering();
+        wallet.executeQuarkOperation(op, signature);
+
+        assertEq(IERC20(WETH).balanceOf(address(wallet)), 10 ether);
+        assertEq(address(wallet).balance, 10 ether);
+    }
+
     function testWrapStETH() public {
         vm.pauseGasMetering();
         QuarkWallet wallet = QuarkWallet(factory.create(alice, address(0)));