test: add unit tests for vault and rewardvault

src/storage/RewardVaultStorage.sol

@@ -15,8 +15,29 @@ contract RewardVaultStorage {
     // Gap for future storage variables
     uint256[40] private _gap;
 
+    /**
+     * @notice Emitted when a reward is deposited.
+     * @param token The address of the token.
+     * @param avs The address of the AVS.
+     * @param amount The amount of the reward deposited.
+     */
     event RewardDeposited(address indexed token, address indexed avs, uint256 amount);
+
+    /**
+     * @notice Emitted when a reward is unlocked.
+     * @param token The address of the token.
+     * @param staker The address of the staker.
+     * @param amount The amount of the reward unlocked.
+     */
     event RewardUnlocked(address indexed token, address indexed staker, uint256 amount);
+
+    /**
+     * @notice Emitted when a reward is withdrawn.
+     * @param token The address of the token.
+     * @param staker The address of the staker.
+     * @param recipient The address of the recipient.
+     * @param amount The amount of the reward withdrawn.
+     */
     event RewardWithdrawn(address indexed token, address indexed staker, address indexed recipient, uint256 amount);
 
 }

src/storage/VaultStorage.sol

@@ -33,18 +33,18 @@ contract VaultStorage {
     /// @notice Emitted when a user's principal balance is deposited.
     /// @param depositor The address of the depositor.
     /// @param amount The amount of the principal balance deposited.
-    event PrincipalDeposited(address depositor, uint256 amount);
+    event PrincipalDeposited(address indexed depositor, uint256 amount);
 
     /// @notice Emitted when a user's principal balance is unlocked for withdrawal.
-    /// @param user The address of the withdrawer.
+    /// @param staker The address of the withdrawer.
     /// @param amount The amount of the principal balance unlocked.
-    event PrincipalUnlocked(address user, uint256 amount);
+    event PrincipalUnlocked(address indexed staker, uint256 amount);
 
     /// @notice Emitted when a user's principal balance is withdrawn.
     /// @param src The address of the withdrawer.
     /// @param dst The address of the recipient.
     /// @param amount The amount of the principal balance withdrawn.
-    event PrincipalWithdrawn(address src, address dst, uint256 amount);
+    event PrincipalWithdrawn(address indexed src, address indexed dst, uint256 amount);
 
     /// @notice Emitted upon the TVL limit being updated.
     /// @param newTvlLimit The new TVL limit.

test/foundry/DepositWithdrawPrinciple.t.sol

@@ -29,9 +29,9 @@ contract DepositWithdrawPrincipalTest is ExocoreDeployer {
     event Transfer(address indexed from, address indexed to, uint256 amount);
     event CapsuleCreated(address owner, address capsule);
     event StakedWithCapsule(address staker, address capsule);
-    event PrincipalDeposited(address depositor, uint256 amount);
-    event PrincipalUnlocked(address user, uint256 amount);
-    event PrincipalWithdrawn(address src, address dst, uint256 amount);
+    event PrincipalDeposited(address indexed depositor, uint256 amount);
+    event PrincipalUnlocked(address indexed staker, uint256 amount);
+    event PrincipalWithdrawn(address indexed src, address indexed dst, uint256 amount);
 
     uint256 constant DEFAULT_ENDPOINT_CALL_GAS_LIMIT = 200_000;
     uint64 public constant MAX_RESTAKED_BALANCE_GWEI_PER_VALIDATOR = 32e9;

test/foundry/unit/RewardVault.t.sol

@@ -0,0 +1,129 @@
+pragma solidity ^0.8.19;
+
+import "../../../src/core/RewardVault.sol";
+
+import "@openzeppelin/contracts/proxy/transparent/ProxyAdmin.sol";
+import "@openzeppelin/contracts/proxy/transparent/TransparentUpgradeableProxy.sol";
+import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
+import "forge-std/Test.sol";
+
+contract MockERC20 is ERC20 {
+
+    constructor(string memory name, string memory symbol) ERC20(name, symbol) {
+        _mint(msg.sender, 1_000_000 * 10 ** 18);
+    }
+
+}
+
+contract RewardVaultTest is Test {
+
+    RewardVault public rewardVaultImplementation;
+    RewardVault public rewardVault;
+    MockERC20 public token;
+    address public gateway;
+    address public depositor;
+    address public withdrawer;
+    address public avs;
+    ProxyAdmin public proxyAdmin;
+
+    event RewardDeposited(address indexed token, address indexed avs, uint256 amount);
+    event RewardWithdrawn(address indexed token, address indexed withdrawer, address indexed recipient, uint256 amount);
+    event RewardUnlocked(address indexed token, address indexed withdrawer, uint256 amount);
+
+    function setUp() public {
+        rewardVaultImplementation = new RewardVault();
+        proxyAdmin = new ProxyAdmin();
+        gateway = address(this);
+
+        // Deploy the proxy
+        TransparentUpgradeableProxy proxy = new TransparentUpgradeableProxy(
+            address(rewardVaultImplementation),
+            address(proxyAdmin),
+            abi.encodeWithSelector(RewardVault.initialize.selector, gateway)
+        );
+
+        // Cast the proxy to RewardVault
+        rewardVault = RewardVault(address(proxy));
+
+        token = new MockERC20("Test Token", "TEST");
+        depositor = address(0x1);
+        withdrawer = address(0x2);
+        avs = address(0x3);
+
+        token.transfer(depositor, 1000 * 10 ** 18);
+    }
+
+    function testInitialize() public {
+        assertEq(rewardVault.gateway(), gateway);
+    }
+
+    function testDeposit() public {
+        uint256 amount = 100 * 10 ** 18;
+        vm.startPrank(depositor);
+        token.approve(address(rewardVault), amount);
+        vm.stopPrank();
+
+        vm.expectEmit(true, true, false, true);
+        emit RewardDeposited(address(token), avs, amount);
+
+        rewardVault.deposit(address(token), depositor, avs, amount);
+
+        assertEq(token.balanceOf(address(rewardVault)), amount);
+        assertEq(rewardVault.getTotalDepositedRewards(address(token), avs), amount);
+    }
+
+    function testWithdraw() public {
+        uint256 amount = 100 * 10 ** 18;
+        rewardVault.unlockReward(address(token), withdrawer, amount);
+        token.transfer(address(rewardVault), amount);
+
+        vm.expectEmit(true, true, true, true);
+        emit RewardWithdrawn(address(token), withdrawer, withdrawer, amount);
+
+        rewardVault.withdraw(address(token), withdrawer, withdrawer, amount);
+
+        assertEq(token.balanceOf(withdrawer), amount);
+        assertEq(rewardVault.getWithdrawableBalance(address(token), withdrawer), 0);
+    }
+
+    function testUnlockReward() public {
+        uint256 amount = 100 * 10 ** 18;
+
+        vm.expectEmit(true, true, false, true);
+        emit RewardUnlocked(address(token), withdrawer, amount);
+
+        rewardVault.unlockReward(address(token), withdrawer, amount);
+
+        assertEq(rewardVault.getWithdrawableBalance(address(token), withdrawer), amount);
+    }
+
+    function testGetWithdrawableBalance() public {
+        uint256 amount = 100 * 10 ** 18;
+        rewardVault.unlockReward(address(token), withdrawer, amount);
+
+        assertEq(rewardVault.getWithdrawableBalance(address(token), withdrawer), amount);
+    }
+
+    function testGetTotalDepositedRewards() public {
+        uint256 amount = 100 * 10 ** 18;
+        vm.startPrank(depositor);
+        token.approve(address(rewardVault), amount);
+        vm.stopPrank();
+
+        rewardVault.deposit(address(token), depositor, avs, amount);
+
+        assertEq(rewardVault.getTotalDepositedRewards(address(token), avs), amount);
+    }
+
+    function testOnlyGatewayModifier() public {
+        vm.prank(address(0x4));
+        vm.expectRevert(Errors.VaultCallerIsNotGateway.selector);
+        rewardVault.deposit(address(token), depositor, avs, 100 * 10 ** 18);
+    }
+
+    function testWithdrawInsufficientBalance() public {
+        vm.expectRevert(Errors.InsufficientBalance.selector);
+        rewardVault.withdraw(address(token), withdrawer, withdrawer, 100 * 10 ** 18);
+    }
+
+}

test/foundry/unit/Vault.t.sol

@@ -0,0 +1,196 @@
+pragma solidity ^0.8.19;
+
+import "../../../src/core/Vault.sol";
+
+import "@openzeppelin/contracts/proxy/transparent/ProxyAdmin.sol";
+import "@openzeppelin/contracts/proxy/transparent/TransparentUpgradeableProxy.sol";
+import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
+import "forge-std/Test.sol";
+
+contract MockERC20 is ERC20 {
+
+    constructor(string memory name, string memory symbol) ERC20(name, symbol) {
+        _mint(msg.sender, 1e10 * 10 ** 18);
+    }
+
+}
+
+contract MockGateway {
+
+    function deposit(address token, address depositor, uint256 amount) external {}
+    function withdraw(address token, address withdrawer, address recipient, uint256 amount) external {}
+
+}
+
+contract VaultTest is Test {
+
+    Vault public vaultImplementation;
+    Vault public vault;
+    MockERC20 public token;
+    MockGateway public gateway;
+    ProxyAdmin public proxyAdmin;
+    address public depositor;
+    address public withdrawer;
+    uint256 public constant TVL_LIMIT = 1_000_000 * 10 ** 18;
+
+    event ConsumedTvlChanged(uint256 newConsumedTvl);
+    event PrincipalDeposited(address indexed depositor, uint256 amount);
+    event PrincipalWithdrawn(address indexed withdrawer, address indexed recipient, uint256 amount);
+    event PrincipalUnlocked(address indexed user, uint256 amount);
+    event TvlLimitUpdated(uint256 newTvlLimit);
+
+    function setUp() public {
+        vaultImplementation = new Vault();
+        proxyAdmin = new ProxyAdmin();
+        token = new MockERC20("Test Token", "TEST");
+        gateway = new MockGateway();
+        depositor = address(0x1);
+        withdrawer = address(0x2);
+
+        // Deploy the proxy
+        TransparentUpgradeableProxy proxy = new TransparentUpgradeableProxy(
+            address(vaultImplementation),
+            address(proxyAdmin),
+            abi.encodeWithSelector(Vault.initialize.selector, address(token), TVL_LIMIT, address(gateway))
+        );
+
+        // Cast the proxy to Vault
+        vault = Vault(address(proxy));
+
+        token.transfer(depositor, TVL_LIMIT + 1000 * 10 ** 18); // Give enough tokens to exceed TVL
+    }
+
+    function testInitialize() public {
+        assertEq(vault.getUnderlyingToken(), address(token));
+        assertEq(vault.getTvlLimit(), TVL_LIMIT);
+        assertEq(address(vault.gateway()), address(gateway));
+    }
+
+    function testDeposit() public {
+        uint256 amount = 100 * 10 ** 18;
+        vm.startPrank(depositor);
+        token.approve(address(vault), amount);
+        vm.stopPrank();
+
+        vm.startPrank(address(gateway));
+        vm.expectEmit(false, false, false, true);
+        emit ConsumedTvlChanged(amount);
+        vm.expectEmit(true, false, false, true);
+        emit PrincipalDeposited(depositor, amount);
+        vault.deposit(depositor, amount);
+        vm.stopPrank();
+
+        assertEq(token.balanceOf(address(vault)), amount);
+        assertEq(vault.totalDepositedPrincipalAmount(depositor), amount);
+        assertEq(vault.getConsumedTvl(), amount);
+    }
+
+    function testWithdraw() public {
+        uint256 amount = 100 * 10 ** 18;
+
+        token.transfer(withdrawer, amount); // Give enough tokens to exceed TVL
+
+        vm.startPrank(withdrawer);
+        token.approve(address(vault), amount);
+        vm.stopPrank();
+
+        vm.startPrank(address(gateway));
+        vault.deposit(withdrawer, amount);
+        vm.stopPrank();
+
+        vm.startPrank(address(gateway));
+        vault.unlockPrincipal(withdrawer, amount);
+        vm.stopPrank();
+
+        vm.startPrank(address(gateway));
+        vm.expectEmit(false, false, false, true);
+        emit ConsumedTvlChanged(0);
+        vm.expectEmit(true, true, false, true);
+        emit PrincipalWithdrawn(withdrawer, withdrawer, amount);
+        vault.withdraw(withdrawer, withdrawer, amount);
+        vm.stopPrank();
+
+        assertEq(token.balanceOf(withdrawer), amount);
+        assertEq(vault.getWithdrawableBalance(withdrawer), 0);
+        assertEq(vault.getConsumedTvl(), 0);
+    }
+
+    function testUnlockPrincipal() public {
+        uint256 amount = 100 * 10 ** 18;
+        vm.startPrank(depositor);
+        token.approve(address(vault), amount);
+        vm.stopPrank();
+
+        vm.startPrank(address(gateway));
+        vault.deposit(depositor, amount);
+        vm.expectEmit(true, false, false, true);
+        emit PrincipalUnlocked(depositor, amount);
+        vault.unlockPrincipal(depositor, amount);
+        vm.stopPrank();
+
+        assertEq(vault.getWithdrawableBalance(depositor), amount);
+    }
+
+    function testSetTvlLimit() public {
+        uint256 newLimit = 2_000_000 * 10 ** 18;
+        vm.startPrank(address(gateway));
+        vm.expectEmit(false, false, false, true);
+        emit TvlLimitUpdated(newLimit);
+        vault.setTvlLimit(newLimit);
+        vm.stopPrank();
+
+        assertEq(vault.getTvlLimit(), newLimit);
+    }
+
+    function testOnlyGatewayModifier() public {
+        vm.expectRevert(Errors.VaultCallerIsNotGateway.selector);
+        vault.deposit(depositor, 100 * 10 ** 18);
+    }
+
+    function testDepositExceedsTvlLimit() public {
+        uint256 amount = TVL_LIMIT + 1;
+        vm.startPrank(depositor);
+        token.approve(address(vault), amount);
+        vm.stopPrank();
+
+        vm.startPrank(address(gateway));
+        vm.expectRevert(Errors.VaultTvlLimitExceeded.selector);
+        vault.deposit(depositor, amount);
+        vm.stopPrank();
+    }
+
+    function testWithdrawExceedsBalance() public {
+        vm.startPrank(address(gateway));
+        vm.expectRevert(Errors.VaultWithdrawalAmountExceeds.selector);
+        vault.withdraw(withdrawer, withdrawer, 1);
+        vm.stopPrank();
+    }
+
+    function testUnlockPrincipalExceedsTotalDeposit() public {
+        uint256 amount = 100 * 10 ** 18;
+        vm.startPrank(depositor);
+        token.approve(address(vault), amount);
+        vm.stopPrank();
+
+        vm.startPrank(address(gateway));
+        vault.deposit(depositor, amount);
+        vm.expectRevert(Errors.VaultPrincipalExceedsTotalDeposit.selector);
+        vault.unlockPrincipal(depositor, amount + 1);
+        vm.stopPrank();
+    }
+
+    function testTotalUnlockPrincipalExceedsDeposit() public {
+        uint256 amount = 100 * 10 ** 18;
+        vm.startPrank(depositor);
+        token.approve(address(vault), amount);
+        vm.stopPrank();
+
+        vm.startPrank(address(gateway));
+        vault.deposit(depositor, amount);
+        vault.unlockPrincipal(depositor, amount / 2);
+        vm.expectRevert(Errors.VaultTotalUnlockPrincipalExceedsDeposit.selector);
+        vault.unlockPrincipal(depositor, (amount / 2) + 1);
+        vm.stopPrank();
+    }
+
+}