From 31997a82672f41e339abe8da76617f9ffdae7852 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=CF=87=C2=B2?=
 <88190723+ChiTimesChi@users.noreply.github.com>
Date: Fri, 22 Nov 2024 19:16:32 +0000
Subject: [PATCH] test(contracts-rfq): gas bench for `TokenZapV1` (#3406)

* test: add minimal mock for a vault contract

* test: gas bench for TokenZap
---
 .../test/mocks/SimpleVaultMock.sol            | 15 +++++
 .../test/mocks/VaultManyArguments.sol         |  4 ++
 .../test/zaps/TokenZapV1.GasBench.t.sol       | 65 +++++++++++++++++++
 3 files changed, 84 insertions(+)
 create mode 100644 packages/contracts-rfq/test/mocks/SimpleVaultMock.sol
 create mode 100644 packages/contracts-rfq/test/zaps/TokenZapV1.GasBench.t.sol

diff --git a/packages/contracts-rfq/test/mocks/SimpleVaultMock.sol b/packages/contracts-rfq/test/mocks/SimpleVaultMock.sol
new file mode 100644
index 0000000000..219976979e
--- /dev/null
+++ b/packages/contracts-rfq/test/mocks/SimpleVaultMock.sol
@@ -0,0 +1,15 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.20;
+
+import {VaultMock} from "./VaultMock.sol";
+
+// solhint-disable no-empty-blocks
+/// @notice Vault mock for testing purposes. DO NOT USE IN PRODUCTION.
+contract SimpleVaultMock is VaultMock {
+    /// @notice We include an empty "test" function so that this contract does not appear in the coverage report.
+    function testSimpleVaultMock() external {}
+
+    function deposit(address token, uint256 amount, address user) external payable {
+        _deposit(user, token, amount);
+    }
+}
diff --git a/packages/contracts-rfq/test/mocks/VaultManyArguments.sol b/packages/contracts-rfq/test/mocks/VaultManyArguments.sol
index 7e43817bd1..20408bf6b9 100644
--- a/packages/contracts-rfq/test/mocks/VaultManyArguments.sol
+++ b/packages/contracts-rfq/test/mocks/VaultManyArguments.sol
@@ -3,10 +3,14 @@ pragma solidity ^0.8.20;
 
 import {VaultMock} from "./VaultMock.sol";
 
+// solhint-disable no-empty-blocks
 /// @notice Vault mock for testing purposes. DO NOT USE IN PRODUCTION.
 contract VaultManyArguments is VaultMock {
     error VaultManyArguments__SomeError();
 
+    /// @notice We include an empty "test" function so that this contract does not appear in the coverage report.
+    function testSimpleVaultMock() external {}
+
     function deposit(
         address token,
         bytes memory encodedToken,
diff --git a/packages/contracts-rfq/test/zaps/TokenZapV1.GasBench.t.sol b/packages/contracts-rfq/test/zaps/TokenZapV1.GasBench.t.sol
new file mode 100644
index 0000000000..bae47da888
--- /dev/null
+++ b/packages/contracts-rfq/test/zaps/TokenZapV1.GasBench.t.sol
@@ -0,0 +1,65 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.24;
+
+import {TokenZapV1} from "../../contracts/zaps/TokenZapV1.sol";
+
+import {SimpleVaultMock} from "../mocks/SimpleVaultMock.sol";
+import {MockERC20} from "../MockERC20.sol";
+
+import {Test} from "forge-std/Test.sol";
+
+// solhint-disable func-name-mixedcase, ordering
+contract TokenZapV1GasBenchmarkTest is Test {
+    uint256 internal constant AMOUNT = 0.1337 ether;
+
+    SimpleVaultMock internal vault;
+    TokenZapV1 internal tokenZap;
+    MockERC20 internal erc20;
+
+    address internal user;
+    address internal nativeGasToken = 0xEeeeeEeeeEeEeeEeEeEeeEEEeeeeEeeeeeeeEEeE;
+
+    function setUp() public {
+        tokenZap = new TokenZapV1();
+        vault = new SimpleVaultMock();
+        erc20 = new MockERC20("TKN", 18);
+
+        user = makeAddr("user");
+
+        erc20.mint(address(this), AMOUNT);
+        deal(address(this), AMOUNT);
+        // To simulate an average case we assume that the Vault contract has already other deposited funds.
+        erc20.mint(address(vault), 1000 * AMOUNT);
+        deal(address(vault), 1000 * AMOUNT);
+        // We also assume that this is not the first tx through the Zap, so the infinite approval has already been set.
+        vm.prank(address(tokenZap));
+        erc20.approve(address(vault), type(uint256).max);
+    }
+
+    function getVaultPayload(address token, uint256 amount) public view returns (bytes memory) {
+        return abi.encodeCall(vault.deposit, (token, amount, user));
+    }
+
+    function getZapData(bytes memory originalPayload) public view returns (bytes memory) {
+        // Amount is the second argument of the deposit function.
+        return tokenZap.encodeZapData(address(vault), originalPayload, 4 + 32);
+    }
+
+    function test_deposit_erc20() public {
+        bytes memory depositPayload = getVaultPayload(address(erc20), AMOUNT);
+        bytes memory zapData = getZapData(depositPayload);
+        // Transfer tokens to the zap contract first.
+        erc20.transfer(address(tokenZap), AMOUNT);
+        tokenZap.zap(address(erc20), AMOUNT, zapData);
+        // Check that the vault registered the deposit.
+        assertEq(vault.balanceOf(user, address(erc20)), AMOUNT);
+    }
+
+    function test_deposit_native() public {
+        bytes memory depositPayload = getVaultPayload(nativeGasToken, AMOUNT);
+        bytes memory zapData = getZapData(depositPayload);
+        tokenZap.zap{value: AMOUNT}(nativeGasToken, AMOUNT, zapData);
+        // Check that the vault registered the deposit.
+        assertEq(vault.balanceOf(user, nativeGasToken), AMOUNT);
+    }
+}