Remove amount from voucher drain

iExecBlockchainComputing · Jun 10, 2024 · e3b20cd · e3b20cd
1 parent fb17fbe
commit e3b20cd
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 10 deletions.
diff --git a/contracts/VoucherHub.sol b/contracts/VoucherHub.sol
@@ -263,7 +263,7 @@ contract VoucherHub is
         require(amount > 0, "VoucherHub: nothing to drain");
         _burn(voucher, amount);
         emit VoucherDrained(voucher, amount);
-        Voucher(voucher).drain(amount);
+        Voucher(voucher).drain();
     }
 
     /**

diff --git a/contracts/beacon/IVoucher.sol b/contracts/beacon/IVoucher.sol
@@ -30,7 +30,7 @@ interface IVoucher {
     ) external returns (bytes32);
     function claim(bytes32 taskId) external;
     function claimBoost(bytes32 dealId, uint256 taskIndex) external;
-    function drain(uint256 amount) external;
+    function drain() external;
 
     function getVoucherHub() external view returns (address);
     function getType() external view returns (uint256);

diff --git a/contracts/beacon/Voucher.sol b/contracts/beacon/Voucher.sol
@@ -271,14 +271,14 @@ contract Voucher is Initializable, IVoucher {
     /**
      * Drain balance of voucher on PoCo if it is expired.
      * Funds are sent to the VoucherHub contract.
-     * @param amount amount to be drained
      */
-    function drain(uint256 amount) external onlyVoucherHub onlyExpired {
+    function drain() external onlyVoucherHub onlyExpired {
         VoucherStorage storage $ = _getVoucherStorage();
         address voucherHub = $._voucherHub;
         // Although transfer function in PoCo always returns true (or reverts),
         // a return value check is added here in case its behavior changes.
-        if (!IERC20(IVoucherHub(voucherHub).getIexecPoco()).transfer(voucherHub, amount)) {
+        IERC20 iexecPoco = IERC20(IVoucherHub(voucherHub).getIexecPoco());
+        if (!iexecPoco.transfer(voucherHub, iexecPoco.balanceOf(address(this)))) {
             revert("Voucher: drain failed");
         }
     }

diff --git a/test/beacon/Voucher.test.ts b/test/beacon/Voucher.test.ts
@@ -1062,23 +1062,23 @@ describe('Voucher', function () {
             await time.setNextBlockTimestamp(expirationDate + 100n); // after expiration
             // Drain
             const voucherHubSigner = await ethers.getImpersonatedSigner(voucherHubAddress);
-            await expect(voucherAsAnyone.connect(voucherHubSigner).drain(voucherValue))
+            await expect(voucherAsAnyone.connect(voucherHubSigner).drain())
                 .to.emit(iexecPocoInstance, 'Transfer')
                 .withArgs(voucherAddress, voucherHubAddress, voucherValue);
             expect(await iexecPocoInstance.balanceOf(voucherAddress)).to.equal(0);
         });
 
         it('Should not drain voucher if sender is not authorized', async function () {
-            await expect(voucherAsAnyone.drain(voucherValue)).to.be.revertedWith(
+            await expect(voucherAsAnyone.drain()).to.be.revertedWith(
                 'Voucher: sender is not VoucherHub',
             );
         });
 
         it('Should not drain voucher if not expired', async function () {
             const voucherHubSigner = await ethers.getImpersonatedSigner(voucherHubAddress);
-            await expect(
-                voucherAsAnyone.connect(voucherHubSigner).drain(voucherValue),
-            ).to.be.revertedWith('Voucher: voucher is not expired');
+            await expect(voucherAsAnyone.connect(voucherHubSigner).drain()).to.be.revertedWith(
+                'Voucher: voucher is not expired',
+            );
         });
     });