diff --git a/src/FoldCaptiveStaking.sol b/src/FoldCaptiveStaking.sol index fdbfec0..7b0ddab 100644 --- a/src/FoldCaptiveStaking.sol +++ b/src/FoldCaptiveStaking.sol @@ -43,6 +43,7 @@ contract FoldCaptiveStaking is Owned(msg.sender) { error AlreadyInitialized(); error NotInitialized(); error ZeroLiquidity(); + error NotEnoughLiquidity(); error WithdrawFailed(); error DepositCapReached(); error DepositAmountBelowMinimum(); @@ -60,13 +61,11 @@ contract FoldCaptiveStaking is Owned(msg.sender) { positionManager = INonfungiblePositionManager(_positionManager); POOL = IUniswapV3Pool(_pool); - token0 = ERC20(POOL.token0()); - token1 = ERC20(POOL.token1()); + token0 = ERC20(IUniswapV3Pool(_pool).token0()); + token1 = ERC20(IUniswapV3Pool(_pool).token1()); WETH9 = WETH(payable(_weth)); FOLD = ERC20(_fold); - - initialized = false; } /// @notice Initialize the contract by minting a small initial liquidity position @@ -293,6 +292,7 @@ contract FoldCaptiveStaking is Owned(msg.sender) { /// @param liquidity The amount of liquidity to withdraw function withdraw(uint128 liquidity) external isInitialized { if (block.timestamp < depositTimeStamp[msg.sender] + COOLDOWN_PERIOD) revert WithdrawalCooldownPeriodNotMet(); + if (liquidity > liquidityUnderManagement) revert NotEnoughLiquidity(); collectFees(); collectRewards(); diff --git a/test/UnitTests.t.sol b/test/UnitTests.t.sol index deeddb7..c7f2e5c 100644 --- a/test/UnitTests.t.sol +++ b/test/UnitTests.t.sol @@ -6,6 +6,19 @@ import "test/interfaces/ISwapRouter.sol"; contract UnitTests is BaseCaptiveTest { ISwapRouter public router = ISwapRouter(0xE592427A0AEce92De3Edee1F18E0157C05861564); + /// @dev Test the constructor that you cannot supply a zero address + function testNoZeroAddressInitialization() public { + vm.expectRevert(FoldCaptiveStaking.ZeroAddress.selector); + new FoldCaptiveStaking(address(0), address(0), address(0), address(0)); + } + + /// @dev Test that you cannot initialize the same contract twice + function testNoDoubleInitialization() public { + vm.startPrank(foldCaptiveStaking.owner()); + vm.expectRevert(FoldCaptiveStaking.AlreadyInitialized.selector); + foldCaptiveStaking.initialize(); + } + /// @dev Ensure that balances and state variables are updated correctly. function testAddLiquidity() public { fold.transfer(User01, 1_000 ether);