Merge pull request #30 from bobanetwork/wsdt/withdraw-partially

feat: Add partial withdraw to contract
bobanetwork · Apr 29, 2024 · 6ae98fd · 6ae98fd
2 parents 8fda290 + 55934c9
commit 6ae98fd
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 16 deletions.
diff --git a/contracts/LightBridge.sol b/contracts/LightBridge.sol
@@ -392,24 +392,24 @@ contract LightBridge is PausableUpgradeable, MulticallUpgradeable {
     /**
      * @dev Sends the contract's current balance to the owner.
      */
-    function withdrawBalance(address _token)
+    function withdrawBalance(address _token, uint256 _amount)
     external
     onlyOwner()
     onlyInitialized()
     {
+        require(_amount != 0, "Amount is 0");
         if (address(0) == _token) {
             uint256 _balance = address(this).balance;
-            require(_balance > 0, "Nothing to send");
-            (bool sent,) = owner.call{gas: 2300, value: _balance}("");
+            require(_balance >= _amount, "Too high");
+            (bool sent,) = owner.call{gas: 2300, value: _amount}("");
             require(sent, "Failed to send Ether");
-            emit AssetBalanceWithdrawn(_token, owner, _balance);
         } else {
             // no supportedToken check in case of generally lost tokens
             uint256 _balance = IERC20(_token).balanceOf(address(this));
-            require(_balance > 0, "Nothing to send");
-            IERC20(_token).safeTransfer(owner, _balance);
-            emit AssetBalanceWithdrawn(_token, owner, _balance);
+            require(_balance >= _amount, "Too high");
+            IERC20(_token).safeTransfer(owner, _amount);
         }
+        emit AssetBalanceWithdrawn(_token, owner, _amount);
     }
 
     /**

diff --git a/test/lightbridge.unit.spec.ts b/test/lightbridge.unit.spec.ts
@@ -596,9 +596,20 @@ describe('Asset Teleportation Tests', async () => {
       it('should withdraw ERC20 balance', async () => {
         const preSignerBalnce = await L2Boba.balanceOf(signerAddress)
         const preBalance = await L2Boba.balanceOf(Proxy__Teleportation.address)
-        await expect(Proxy__Teleportation.withdrawBalance(L2Boba.address))
+
+        const partialAmount = preBalance.div(3)
+        await expect(
+          Proxy__Teleportation.withdrawBalance(L2Boba.address, partialAmount)
+        )
+          .to.emit(Proxy__Teleportation, 'AssetBalanceWithdrawn')
+          .withArgs(L2Boba.address, signerAddress, partialAmount)
+
+        const newPreBalance = await L2Boba.balanceOf(Proxy__Teleportation.address)
+        await expect(
+          Proxy__Teleportation.withdrawBalance(L2Boba.address, newPreBalance)
+        )
           .to.emit(Proxy__Teleportation, 'AssetBalanceWithdrawn')
-          .withArgs(L2Boba.address, signerAddress, preBalance)
+          .withArgs(L2Boba.address, signerAddress, newPreBalance)
 
         const postBalance = await L2Boba.balanceOf(Proxy__Teleportation.address)
         const postSignerBalance = await L2Boba.balanceOf(signerAddress)
@@ -610,7 +621,10 @@ describe('Asset Teleportation Tests', async () => {
 
       it('should not withdraw ERC20 balance if caller is not owner', async () => {
         await expect(
-          Proxy__Teleportation.connect(signer2).withdrawBalance(L2Boba.address)
+          Proxy__Teleportation.connect(signer2).withdrawBalance(
+            L2Boba.address,
+            '1'
+          )
         ).to.be.revertedWith('Caller is not the owner')
       })
 
@@ -1249,15 +1263,32 @@ describe('Asset Teleportation Tests', async () => {
       })
 
       it('should withdraw BOBA balance', async () => {
-        const preSignerBalnce = await ethers.provider.getBalance(signerAddress)
+        const preSignerBalance = await ethers.provider.getBalance(signerAddress)
         const preBalance = await ethers.provider.getBalance(
           Proxy__Teleportation.address
         )
+
+        const partialAmount = preBalance.div(3)
+        await expect(
+          Proxy__Teleportation.withdrawBalance(
+            ethers.constants.AddressZero,
+            partialAmount
+          )
+        )
+          .to.emit(Proxy__Teleportation, 'AssetBalanceWithdrawn')
+          .withArgs(ethers.constants.AddressZero, signerAddress, partialAmount)
+
+        const newPreBalance = await ethers.provider.getBalance(
+          Proxy__Teleportation.address
+        )
         await expect(
-          Proxy__Teleportation.withdrawBalance(ethers.constants.AddressZero)
+          Proxy__Teleportation.withdrawBalance(
+            ethers.constants.AddressZero,
+            newPreBalance
+          )
         )
           .to.emit(Proxy__Teleportation, 'AssetBalanceWithdrawn')
-          .withArgs(ethers.constants.AddressZero, signerAddress, preBalance)
+          .withArgs(ethers.constants.AddressZero, signerAddress, newPreBalance)
 
         const postBalance = await ethers.provider.getBalance(
           Proxy__Teleportation.address
@@ -1266,16 +1297,18 @@ describe('Asset Teleportation Tests', async () => {
           signerAddress
         )
         const gasFee = await getGasFeeFromLastestBlock(ethers.provider)
-        expect(preBalance.sub(postBalance)).to.be.eq(
-          postSignerBalance.sub(preSignerBalnce).add(gasFee)
+        expect(preBalance.sub(postBalance)).to.be.closeTo(
+          postSignerBalance.sub(preSignerBalance),
+          postSignerBalance.sub(preSignerBalance).add(gasFee),
         )
         expect(postBalance.toString()).to.be.eq('0')
       })
 
       it('should not withdraw BOBA balance if caller is not owner', async () => {
         await expect(
           Proxy__Teleportation.connect(signer2).withdrawBalance(
-            ethers.constants.AddressZero
+            ethers.constants.AddressZero,
+            '1'
           )
         ).to.be.revertedWith('Caller is not the owner')
       })