fix bugs change tests

contracts/Airdrop.sol

@@ -1,24 +1,25 @@
 // SPDX-License-Identifier: UNLICENSED
-pragma solidity =0.5.17;
+pragma solidity ^0.8.9;
-
-pragma experimental ABIEncoderV2;
 
 import "./interfaces/IERC20.sol";
 import "./interfaces/IGovernance.sol";
-import "./libraries/SafeERC20.sol";
 import "./ReentrancyGuard.sol";
-import "./CarefulMath.sol";
 
 import "./interfaces/ISablierAirdrop.sol";
 import "./interfaces/IRecipientStorage.sol";
 import "./libraries/Types.sol";
 
-contract SablierAirdrop is ISablierAirdrop, ReentrancyGuard, CarefulMath {
+contract SablierAirdrop is ISablierAirdrop, ReentrancyGuard {
-    using SafeERC20 for IERC20;
-
     /*** Storage Properties ***/
 
+    /**
+     * @notice Tornado Governance contract is an owner of airdrop contract and can create or cancel airdrops
+     */
     address public constant tornadoGovernance = 0x5efda50f22d34F262c29268506C5Fa42cB56A1Ce;
+
+    /**
+     * @notice TORN token - token for airdrops
+     */
     IERC20 public constant torn = IERC20(0x77777FeDdddFfC19Ff86DB637967013e6C6A116C);
 
     /**
@@ -59,7 +60,7 @@ contract SablierAirdrop is ISablierAirdrop, ReentrancyGuard, CarefulMath {
 
     /*** Contract Logic Starts Here */
 
-    constructor() public {
+    constructor() {
         nextStreamId = 1;
     }
 
@@ -69,7 +70,6 @@ contract SablierAirdrop is ISablierAirdrop, ReentrancyGuard, CarefulMath {
      * @notice Returns the stream with all its properties.
      * @dev Throws if the id does not point to a valid stream.
      * @param streamId The id of the stream to query.
-     * @return The stream object.
      */
     function getStream(
         uint256 streamId
@@ -100,7 +100,7 @@ contract SablierAirdrop is ISablierAirdrop, ReentrancyGuard, CarefulMath {
      *  `startTime`, it returns 0.
      * @dev Throws if the id does not point to a valid stream.
      * @param streamId The id of the stream for which to query the delta.
-     * @return The time delta in seconds.
+     * @return delta The time delta in seconds.
      */
     function deltaOf(uint256 streamId) public view streamExists(streamId) returns (uint256 delta) {
         Types.Stream memory stream = streams[streamId];
@@ -109,26 +109,17 @@ contract SablierAirdrop is ISablierAirdrop, ReentrancyGuard, CarefulMath {
         return stream.stopTime - stream.startTime;
     }
 
-    struct BalanceOfLocalVars {
-        MathError mathErr;
-        uint256 recipientBalance;
-        uint256 withdrawalAmount;
-    }
-
     /**
      * @notice Returns the available funds for the given stream id and address.
      * @dev Throws if the id does not point to a valid stream.
      * @param streamId The id of the stream for which to query the balance.
      * @param who The address for which to query the balance.
-     * @return The total funds allocated to `who` as uint256.
+     * @return balance The total funds allocated to `who` as uint256.
      */
     function balanceOf(uint256 streamId, address who) public view streamExists(streamId) returns (uint256 balance) {
         Types.Stream memory stream = streams[streamId];
-        BalanceOfLocalVars memory vars;
-
         uint256 delta = deltaOf(streamId);
-        (vars.mathErr, vars.recipientBalance) = mulUInt(delta, stream.ratePerSecond);
+        uint256 recipientBalance = delta * stream.ratePerSecond;
-        require(vars.mathErr == MathError.NO_ERROR, "recipient balance calculation error");
 
         /*
          * If the stream `balance` does not equal `deposit`, it means there have been withdrawals.
@@ -136,14 +127,11 @@ contract SablierAirdrop is ISablierAirdrop, ReentrancyGuard, CarefulMath {
          * streamed until now.
          */
         if (stream.deposit > stream.remainingBalance) {
-            (vars.mathErr, vars.withdrawalAmount) = subUInt(stream.deposit, stream.remainingBalance);
+            uint256 withdrawalAmount = stream.deposit - stream.remainingBalance;
-            assert(vars.mathErr == MathError.NO_ERROR);
+            recipientBalance = recipientBalance - withdrawalAmount;
-            (vars.mathErr, vars.recipientBalance) = subUInt(vars.recipientBalance, vars.withdrawalAmount);
-            /* `withdrawalAmount` cannot and should not be bigger than `recipientBalance`. */
-            assert(vars.mathErr == MathError.NO_ERROR);
         }
 
-        if (who == stream.recipient) return vars.recipientBalance;
+        if (who == stream.recipient) return recipientBalance;
         return 0;
     }
 
@@ -151,7 +139,6 @@ contract SablierAirdrop is ISablierAirdrop, ReentrancyGuard, CarefulMath {
      * @notice To avoid "stack to deep" error
      */
     struct CreateAirdropLocalVars {
-        MathError mathErr;
         uint256 airdropDuration;
         uint256 ratePerSecond;
         uint256 firstStream;
@@ -163,7 +150,7 @@ contract SablierAirdrop is ISablierAirdrop, ReentrancyGuard, CarefulMath {
         uint256 startTime,
         uint256 stopTime,
         address recipientStorage
-    ) public onlyGovernance returns (bool) {
+    ) external onlyGovernance returns (bool) {
         CreateAirdropLocalVars memory vars;
         vars.airdropDuration = stopTime - startTime;
         vars.airdropRecipients = IRecipientStorage(recipientStorage).getAirdropRecipients();
@@ -180,9 +167,7 @@ contract SablierAirdrop is ISablierAirdrop, ReentrancyGuard, CarefulMath {
             /* Without this, the rate per second would be zero. */
             require(normalizedDeposit >= vars.airdropDuration, "deposit smaller than time delta");
 
-            (vars.mathErr, vars.ratePerSecond) = divUInt(normalizedDeposit, vars.airdropDuration);
+            vars.ratePerSecond = normalizedDeposit / vars.airdropDuration;
-            /* `divUInt` can only return MathError.DIVISION_BY_ZERO but we know `duration` is not zero. */
-            assert(vars.mathErr == MathError.NO_ERROR);
 
             /* Create and store the stream object. */
             uint256 streamId = nextStreamId;
@@ -198,8 +183,7 @@ contract SablierAirdrop is ISablierAirdrop, ReentrancyGuard, CarefulMath {
             });
 
             /* Increment the next stream id. */
-            (vars.mathErr, nextStreamId) = addUInt(nextStreamId, uint256(1));
+            nextStreamId = nextStreamId + 1;
-            require(vars.mathErr == MathError.NO_ERROR, "next stream id calculation error");
 
             emit CreateStream(
                 streamId,
@@ -216,39 +200,31 @@ contract SablierAirdrop is ISablierAirdrop, ReentrancyGuard, CarefulMath {
     }
 
     /**
-     * @notice Withdraws from the contract to the recipient's account.
+     * @notice Withdraws from the contract to the recipient's account all available recipient stream balance.
      * @dev Throws if the id does not point to a valid stream.
      *  Throws if the caller is not the sender or the recipient of the stream.
      *  Throws if the amount exceeds the available balance.
      *  Throws if there is a token transfer failure.
      * @param streamId The id of the stream to withdraw tokens from.
-     * @param amount The amount of tokens to withdraw.
      */
     function withdrawFromStream(
-        uint256 streamId,
+        uint256 streamId
-        uint256 amount
     ) external nonReentrant streamExists(streamId) onlyRecipient(streamId) returns (bool) {
-        require(amount > 0, "amount is zero");
         Types.Stream memory stream = streams[streamId];
 
+        uint256 balance = balanceOf(streamId, stream.recipient);
+        require(balance > 0, "amount is zero");
+
         uint256 recipientLockedBalance = IGovernance(tornadoGovernance).lockedBalance(stream.recipient);
         require(recipientLockedBalance >= stream.initialLockedBalance, "not enough locked tokens in governance");
 
-        uint256 balance = balanceOf(streamId, stream.recipient);
+        /* Remaining balance can not be less than recipient stream balance */
-        require(balance >= amount, "amount exceeds the available balance");
+        streams[streamId].remainingBalance = stream.remainingBalance - balance;
-
-        MathError mathErr;
-        (mathErr, streams[streamId].remainingBalance) = subUInt(stream.remainingBalance, amount);
-        /**
-         * `subUInt` can only return MathError.INTEGER_UNDERFLOW but we know that `remainingBalance` is at least
-         * as big as `amount`.
-         */
-        assert(mathErr == MathError.NO_ERROR);
 
         if (streams[streamId].remainingBalance == 0) delete streams[streamId];
 
-        torn.safeTransfer(stream.recipient, amount);
+        torn.transfer(stream.recipient, balance);
-        emit WithdrawFromStream(streamId, stream.recipient, amount);
+        emit WithdrawFromStream(streamId, stream.recipient, balance);
         return true;
     }
 
@@ -266,7 +242,7 @@ contract SablierAirdrop is ISablierAirdrop, ReentrancyGuard, CarefulMath {
 
         delete streams[streamId];
 
-        if (remainingBalance > 0) torn.safeTransfer(tornadoGovernance, remainingBalance);
+        if (remainingBalance > 0) torn.transfer(tornadoGovernance, remainingBalance);
 
         emit CancelStream(streamId, stream.recipient, remainingBalance);
         return true;
@@ -293,7 +269,7 @@ contract SablierAirdrop is ISablierAirdrop, ReentrancyGuard, CarefulMath {
             }
         }
 
-        torn.safeTransfer(tornadoGovernance, airdropRemainingBalance);
+        torn.transfer(tornadoGovernance, airdropRemainingBalance);
         return true;
     }
 
@@ -303,7 +279,7 @@ contract SablierAirdrop is ISablierAirdrop, ReentrancyGuard, CarefulMath {
      * @return bool true=success, otherwise false.
      */
     function withdrawFunds() external onlyGovernance returns (bool) {
-        torn.safeTransfer(tornadoGovernance, torn.balanceOf(address(this)));
+        torn.transfer(tornadoGovernance, torn.balanceOf(address(this)));
         return true;
     }
 }

contracts/CarefulMath.sol

@@ -1,84 +0,0 @@
-pragma solidity >=0.5.17;
-
-/**
- * @title Careful Math
- * @author Compound
- * @notice Derived from OpenZeppelin's SafeMath library
- *         https://github.com/OpenZeppelin/openzeppelin-solidity/blob/master/contracts/math/SafeMath.sol
- */
-contract CarefulMath {
-    /**
-     * @dev Possible error codes that we can return
-     */
-    enum MathError {
-        NO_ERROR,
-        DIVISION_BY_ZERO,
-        INTEGER_OVERFLOW,
-        INTEGER_UNDERFLOW
-    }
-
-    /**
-     * @dev Multiplies two numbers, returns an error on overflow.
-     */
-    function mulUInt(uint a, uint b) internal pure returns (MathError, uint) {
-        if (a == 0) {
-            return (MathError.NO_ERROR, 0);
-        }
-
-        uint c = a * b;
-
-        if (c / a != b) {
-            return (MathError.INTEGER_OVERFLOW, 0);
-        } else {
-            return (MathError.NO_ERROR, c);
-        }
-    }
-
-    /**
-     * @dev Integer division of two numbers, truncating the quotient.
-     */
-    function divUInt(uint a, uint b) internal pure returns (MathError, uint) {
-        if (b == 0) {
-            return (MathError.DIVISION_BY_ZERO, 0);
-        }
-
-        return (MathError.NO_ERROR, a / b);
-    }
-
-    /**
-     * @dev Subtracts two numbers, returns an error on overflow (i.e. if subtrahend is greater than minuend).
-     */
-    function subUInt(uint a, uint b) internal pure returns (MathError, uint) {
-        if (b <= a) {
-            return (MathError.NO_ERROR, a - b);
-        } else {
-            return (MathError.INTEGER_UNDERFLOW, 0);
-        }
-    }
-
-    /**
-     * @dev Adds two numbers, returns an error on overflow.
-     */
-    function addUInt(uint a, uint b) internal pure returns (MathError, uint) {
-        uint c = a + b;
-
-        if (c >= a) {
-            return (MathError.NO_ERROR, c);
-        } else {
-            return (MathError.INTEGER_OVERFLOW, 0);
-        }
-    }
-
-    /**
-     * @dev add a and b and then subtract c
-     */
-    function addThenSubUInt(uint a, uint b, uint c) internal pure returns (MathError, uint) {
-        (MathError err0, uint sum) = addUInt(a, b);
-
-        if (err0 != MathError.NO_ERROR) {
-            return (err0, 0);
-        }
-
-        return subUInt(sum, c);
-    }
-}

contracts/ReentrancyGuard.sol

@@ -1,4 +1,4 @@
-pragma solidity ^0.5.0;
+pragma solidity ^0.8.9;
 
 /**
  * @dev Contract module that helps prevent reentrant calls to a function.
@@ -16,7 +16,7 @@ contract ReentrancyGuard {
     /// @dev counter to allow mutex lock with only one SSTORE operation
     uint256 private _guardCounter;
 
-    constructor() internal {
+    constructor() {
         // The counter starts at one to prevent changing it from zero to a non-zero
         // value, which is a more expensive operation.
         _guardCounter = 1;

contracts/interfaces/IERC20.sol

@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: UNLICENSED
-pragma solidity >=0.5.0;
+pragma solidity ^0.8.9;
 
 /**
  * @dev Interface of the ERC20 standard as defined in the EIP. Does not include

contracts/interfaces/IGovernance.sol

@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: UNLICENSED
-pragma solidity >=0.5.0;
+pragma solidity ^0.8.9;
 
 interface IGovernance {
     function lockedBalance(address staker) external view returns (uint256);

contracts/interfaces/IRecipientStorage.sol

@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: UNLICENSED
-pragma solidity >=0.5.17;
+pragma solidity ^0.8.9;
 
 pragma experimental ABIEncoderV2;
 

contracts/interfaces/ISablierAirdrop.sol

@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: UNLICENSED
-pragma solidity >=0.5.17;
+pragma solidity ^0.8.9;
 
 pragma experimental ABIEncoderV2;
 
@@ -57,7 +57,7 @@ interface ISablierAirdrop {
 
     function createAirdrop(uint256 startTime, uint256 stopTime, address recipientStorage) external returns (bool);
 
-    function withdrawFromStream(uint256 streamId, uint256 funds) external returns (bool);
+    function withdrawFromStream(uint256 streamId) external returns (bool);
 
     function cancelStream(uint256 streamId) external returns (bool);
 }

contracts/libraries/Address.sol

@@ -1,4 +1,4 @@
-pragma solidity ^0.5.0;
+pragma solidity ^0.8.9;
 
 /**
  * @dev Collection of functions related to the address type,

contracts/libraries/SafeERC20.sol

@@ -1,77 +0,0 @@
-pragma solidity ^0.5.0;
-
-import "../interfaces/IERC20.sol";
-import "./SafeMath.sol";
-import "./Address.sol";
-
-/**
- * @title SafeERC20
- * @dev Wrappers around ERC20 operations that throw on failure (when the token
- * contract returns false). Tokens that return no value (and instead revert or
- * throw on failure) are also supported, non-reverting calls are assumed to be
- * successful.
- * To use this library you can add a `using SafeERC20 for ERC20;` statement to your contract,
- * which allows you to call the safe operations as `token.safeTransfer(...)`, etc.
- */
-library SafeERC20 {
-    using SafeMath for uint256;
-    using Address for address;
-
-    function safeTransfer(IERC20 token, address to, uint256 value) internal {
-        callOptionalReturn(token, abi.encodeWithSelector(token.transfer.selector, to, value));
-    }
-
-    function safeTransferFrom(IERC20 token, address from, address to, uint256 value) internal {
-        callOptionalReturn(token, abi.encodeWithSelector(token.transferFrom.selector, from, to, value));
-    }
-
-    function safeApprove(IERC20 token, address spender, uint256 value) internal {
-        // safeApprove should only be called when setting an initial allowance,
-        // or when resetting it to zero. To increase and decrease it, use
-        // 'safeIncreaseAllowance' and 'safeDecreaseAllowance'
-        // solhint-disable-next-line max-line-length
-        require(
-            (value == 0) || (token.allowance(address(this), spender) == 0),
-            "SafeERC20: approve from non-zero to non-zero allowance"
-        );
-        callOptionalReturn(token, abi.encodeWithSelector(token.approve.selector, spender, value));
-    }
-
-    function safeIncreaseAllowance(IERC20 token, address spender, uint256 value) internal {
-        uint256 newAllowance = token.allowance(address(this), spender).add(value);
-        callOptionalReturn(token, abi.encodeWithSelector(token.approve.selector, spender, newAllowance));
-    }
-
-    function safeDecreaseAllowance(IERC20 token, address spender, uint256 value) internal {
-        uint256 newAllowance = token.allowance(address(this), spender).sub(value);
-        callOptionalReturn(token, abi.encodeWithSelector(token.approve.selector, spender, newAllowance));
-    }
-
-    /**
-     * @dev Imitates a Solidity high-level call (i.e. a regular function call to a contract), relaxing the requirement
-     * on the return value: the return value is optional (but if data is returned, it must not be false).
-     * @param token The token targeted by the call.
-     * @param data The call data (encoded using abi.encode or one of its variants).
-     */
-    function callOptionalReturn(IERC20 token, bytes memory data) private {
-        // We need to perform a low level call here, to bypass Solidity's return data size checking mechanism, since
-        // we're implementing it ourselves.
-
-        // A Solidity high level call has three parts:
-        //  1. The target address is checked to verify it contains contract code
-        //  2. The call itself is made, and success asserted
-        //  3. The return value is decoded, which in turn checks the size of the returned data.
-        // solhint-disable-next-line max-line-length
-        require(address(token).isContract(), "SafeERC20: call to non-contract");
-
-        // solhint-disable-next-line avoid-low-level-calls
-        (bool success, bytes memory returndata) = address(token).call(data);
-        require(success, "SafeERC20: low-level call failed");
-
-        if (returndata.length > 0) {
-            // Return data is optional
-            // solhint-disable-next-line max-line-length
-            require(abi.decode(returndata, (bool)), "SafeERC20: ERC20 operation did not succeed");
-        }
-    }
-}

contracts/libraries/SafeMath.sol

@@ -1,107 +0,0 @@
-pragma solidity ^0.5.0;
-
-/**
- * @dev Wrappers over Solidity's arithmetic operations with added overflow
- * checks.
- *
- * Arithmetic operations in Solidity wrap on overflow. This can easily result
- * in bugs, because programmers usually assume that an overflow raises an
- * error, which is the standard behavior in high level programming languages.
- * `SafeMath` restores this intuition by reverting the transaction when an
- * operation overflows.
- *
- * Using this library instead of the unchecked operations eliminates an entire
- * class of bugs, so it's recommended to use it always.
- */
-library SafeMath {
-    /**
-     * @dev Returns the addition of two unsigned integers, reverting on
-     * overflow.
-     *
-     * Counterpart to Solidity's `+` operator.
-     *
-     * Requirements:
-     * - Addition cannot overflow.
-     */
-    function add(uint256 a, uint256 b) internal pure returns (uint256) {
-        uint256 c = a + b;
-        require(c >= a, "SafeMath: addition overflow");
-
-        return c;
-    }
-
-    /**
-     * @dev Returns the subtraction of two unsigned integers, reverting on
-     * overflow (when the result is negative).
-     *
-     * Counterpart to Solidity's `-` operator.
-     *
-     * Requirements:
-     * - Subtraction cannot overflow.
-     */
-    function sub(uint256 a, uint256 b) internal pure returns (uint256) {
-        require(b <= a, "SafeMath: subtraction overflow");
-        uint256 c = a - b;
-
-        return c;
-    }
-
-    /**
-     * @dev Returns the multiplication of two unsigned integers, reverting on
-     * overflow.
-     *
-     * Counterpart to Solidity's `*` operator.
-     *
-     * Requirements:
-     * - Multiplication cannot overflow.
-     */
-    function mul(uint256 a, uint256 b) internal pure returns (uint256) {
-        // Gas optimization: this is cheaper than requiring 'a' not being zero, but the
-        // benefit is lost if 'b' is also tested.
-        // See: https://github.com/OpenZeppelin/openzeppelin-contracts/pull/522
-        if (a == 0) {
-            return 0;
-        }
-
-        uint256 c = a * b;
-        require(c / a == b, "SafeMath: multiplication overflow");
-
-        return c;
-    }
-
-    /**
-     * @dev Returns the integer division of two unsigned integers. Reverts on
-     * division by zero. The result is rounded towards zero.
-     *
-     * Counterpart to Solidity's `/` operator. Note: this function uses a
-     * `revert` opcode (which leaves remaining gas untouched) while Solidity
-     * uses an invalid opcode to revert (consuming all remaining gas).
-     *
-     * Requirements:
-     * - The divisor cannot be zero.
-     */
-    function div(uint256 a, uint256 b) internal pure returns (uint256) {
-        // Solidity only automatically asserts when dividing by 0
-        require(b > 0, "SafeMath: division by zero");
-        uint256 c = a / b;
-        // assert(a == b * c + a % b); // There is no case in which this doesn't hold
-
-        return c;
-    }
-
-    /**
-     * @dev Returns the remainder of dividing two unsigned integers. (unsigned integer modulo),
-     * Reverts when dividing by zero.
-     *
-     * Counterpart to Solidity's `%` operator. This function uses a `revert`
-     * opcode (which leaves remaining gas untouched) while Solidity uses an
-     * invalid opcode to revert (consuming all remaining gas).
-     *
-     * Requirements:
-     * - The divisor cannot be zero.
-     */
-    function mod(uint256 a, uint256 b) internal pure returns (uint256) {
-        require(b != 0, "SafeMath: modulo by zero");
-        return a % b;
-    }
-}

contracts/libraries/Types.sol

@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: UNLICENSED
-pragma solidity >=0.5.17;
+pragma solidity ^0.8.9;
 
 library Types {
     struct Stream {

hardhat.config.js

@@ -8,13 +8,7 @@ module.exports = {
     compilers: [
       {
         version: "0.8.20",
-      },
+      }
-      {
-        version: "0.6.12",
-      },
-      {
-        version: "0.5.17",
-      },
     ],
   },
   settings: {

package-lock.json

@@ -9,7 +9,7 @@
       "version": "1.0.0",
       "license": "ISC",
       "dependencies": {
-        "@openzeppelin/contracts": "^3.2.0-rc.0",
+        "@openzeppelin/contracts": "^3.4.2",
         "@openzeppelin/upgrades-core": "^1.30.1",
         "base58-solidity": "^1.0.2",
         "bignumber.js": "^9.0.1",

package.json

@@ -36,7 +36,7 @@
     "typescript": "^5.2.2"
   },
   "dependencies": {
-    "@openzeppelin/contracts": "^3.2.0-rc.0",
+    "@openzeppelin/contracts": "^3.4.2",
     "@openzeppelin/upgrades-core": "^1.30.1",
     "base58-solidity": "^1.0.2",
     "bignumber.js": "^9.0.1",

test/Proposal.js

@@ -125,37 +125,39 @@ describe("Proposal results check", function () {
     const torn = await getTorn();
     const recipientBalance = await torn.balanceOf(someRecipient);
     const connectedAirdrop = airdropContract.connect(await ethers.getImpersonatedSigner(someRecipient));
-    await connectedAirdrop.withdrawFromStream(recipientStreamId, normalizedAmount);
+    await connectedAirdrop.withdrawFromStream(recipientStreamId);
 
     expect(await torn.balanceOf(someRecipient)).to.be.equal(recipientBalance + normalizedAmount);
   });
 
-  it("Airdrop recipient should not be able to withdraw all funds early", async function () {
+  it("Airdrop recipient should be able to withdraw part of funds early", async function () {
     const someRecipient = await resolveAddr("butterfly-effect.eth");
-    const { airdropContract, airdropRecipientsContract } = await deployAndExecuteProposal();
+    const { airdropContract } = await deployAndExecuteProposal();
-    const recipients = await airdropRecipientsContract.getAirdropRecipients();
-    const selectedRecipientInfo = recipients.find((r) => r[0] === someRecipient);
-    const recipientAirdropAmount = selectedRecipientInfo[1];
     const halfYear = 180 * 24 * 60 * 60;
-    const normalizedAmount = normalizeAirdropAmount(recipientAirdropAmount);
     const events = await getEvents(airdropContract, "CreateStream");
     const recipientStreamId = events.find((e) => e.args[1] === someRecipient).args[0];
 
     await time.increase(halfYear / 2);
+    const torn = await getTorn();
+    const recipientBalance = await torn.balanceOf(someRecipient);
     const connectedAirdrop = airdropContract.connect(await ethers.getImpersonatedSigner(someRecipient));
-    await expect(connectedAirdrop.withdrawFromStream(recipientStreamId, normalizedAmount)).to.be.revertedWith(
+
-      "amount exceeds the available balance",
+    // Add timestamps and manual calculation bcs withdrawFromStream call mine next block and increase stream balance
-    );
+    const streamBalance = await connectedAirdrop.balanceOf(recipientStreamId, someRecipient);
+    const beforeWithdrawalTimestamp = (await ethers.provider.getBlock(await ethers.provider.getBlockNumber())).timestamp;
+    await connectedAirdrop.withdrawFromStream(recipientStreamId);
+    const afterWithdrawalTimestamp = (await ethers.provider.getBlock(await ethers.provider.getBlockNumber())).timestamp;
+    const blockChangeTime = afterWithdrawalTimestamp - beforeWithdrawalTimestamp;
+    const stream = await connectedAirdrop.getStream(recipientStreamId);
+    const ratePerSecond = stream[5];
+    const addedInCurrentBlock = BigInt(blockChangeTime) * ratePerSecond;
+    expect(await torn.balanceOf(someRecipient)).to.be.equal(recipientBalance + streamBalance + addedInCurrentBlock);
   });
 
   it("Airdrop recipient should not be able to withdraw funds if his stake balance is lower than initial", async function () {
     const someRecipient = await resolveAddr("butterfly-effect.eth");
-    const { airdropContract, airdropRecipientsContract } = await deployAndExecuteProposal();
+    const { airdropContract } = await deployAndExecuteProposal();
-    const recipients = await airdropRecipientsContract.getAirdropRecipients();
-    const selectedRecipientInfo = recipients.find((r) => r[0] === someRecipient);
-    const recipientAirdropAmount = selectedRecipientInfo[1];
     const halfYear = 180 * 24 * 60 * 60;
-    const normalizedAmount = normalizeAirdropAmount(recipientAirdropAmount);
     const events = await getEvents(airdropContract, "CreateStream");
     const recipientStreamId = events.find((e) => e.args[1] === someRecipient).args[0];
 
@@ -163,44 +165,22 @@ describe("Proposal results check", function () {
     const governance = await getGovernance(await ethers.getImpersonatedSigner(someRecipient));
     governance.unlock(1000n * 10n ** 18n);
     const connectedAirdrop = airdropContract.connect(await ethers.getImpersonatedSigner(someRecipient));
-    await expect(connectedAirdrop.withdrawFromStream(recipientStreamId, normalizedAmount)).to.be.revertedWith(
+    await expect(connectedAirdrop.withdrawFromStream(recipientStreamId)).to.be.revertedWith(
       "not enough locked tokens in governance",
     );
   });
 
-  it("Airdrop recipient should be able to withdraw part of funds", async function () {
-    const someRecipient = await resolveAddr("butterfly-effect.eth");
-    const { airdropContract, airdropRecipientsContract } = await deployAndExecuteProposal();
-    const recipients = await airdropRecipientsContract.getAirdropRecipients();
-    const selectedRecipientInfo = recipients.find((r) => r[0] === someRecipient);
-    const recipientAirdropAmount = selectedRecipientInfo[1];
-    const quarter = 90 * 24 * 60 * 60;
-    const events = await getEvents(airdropContract, "CreateStream");
-    const recipientStreamId = events.find((e) => e.args[1] === someRecipient).args[0];
-
-    await time.increase(quarter);
-    const torn = await getTorn();
-    const recipientBalance = await torn.balanceOf(someRecipient);
-    const connectedAirdrop = airdropContract.connect(await ethers.getImpersonatedSigner(someRecipient));
-    await connectedAirdrop.withdrawFromStream(recipientStreamId, recipientAirdropAmount / 4n);
-    expect(await torn.balanceOf(someRecipient)).to.be.equal(recipientBalance + recipientAirdropAmount / 4n);
-  });
-
   it("Airdrop recipient should not be able to withdraw funds if stream is canceled", async function () {
     const someRecipient = await resolveAddr("butterfly-effect.eth");
     const { airdropContract, airdropRecipientsContract } = await deployAndExecuteProposal();
-    const recipients = await airdropRecipientsContract.getAirdropRecipients();
-    const selectedRecipientInfo = recipients.find((r) => r[0] === someRecipient);
-    const recipientAirdropAmount = selectedRecipientInfo[1];
     const halfYear = 180 * 24 * 60 * 60;
-    const normalizedAmount = recipientAirdropAmount - (recipientAirdropAmount % BigInt(halfYear));
     const events = await getEvents(airdropContract, "CreateStream");
     const recipientStreamId = events.find((e) => e.args[1] === someRecipient).args[0];
 
     await time.increase(halfYear);
     await airdropContract.cancelStream(recipientStreamId);
     const connectedAirdrop = airdropContract.connect(await ethers.getImpersonatedSigner(someRecipient));
-    await expect(connectedAirdrop.withdrawFromStream(recipientStreamId, normalizedAmount)).to.be.revertedWith(
+    await expect(connectedAirdrop.withdrawFromStream(recipientStreamId)).to.be.revertedWith(
       "stream does not exist",
     );
   });
@@ -209,17 +189,14 @@ describe("Proposal results check", function () {
     const someRecipient = await resolveAddr("butterfly-effect.eth");
     const { airdropContract, airdropRecipientsContract } = await deployAndExecuteProposal();
     const recipients = await airdropRecipientsContract.getAirdropRecipients();
-    const selectedRecipientInfo = recipients.find((r) => r[0] === someRecipient);
-    const recipientAirdropAmount = selectedRecipientInfo[1];
     const halfYear = 180 * 24 * 60 * 60;
-    const normalizedAmount = normalizeAirdropAmount(recipientAirdropAmount);
     const events = await getEvents(airdropContract, "CreateStream");
     const recipientStreamId = events.find((e) => e.args[1] === someRecipient).args[0];
 
     await time.increase(halfYear);
     await airdropContract.cancelAirdrop(1, recipients.length);
     const connectedAirdrop = airdropContract.connect(await ethers.getImpersonatedSigner(someRecipient));
-    await expect(connectedAirdrop.withdrawFromStream(recipientStreamId, normalizedAmount)).to.be.revertedWith(
+    await expect(connectedAirdrop.withdrawFromStream(recipientStreamId)).to.be.revertedWith(
       "stream does not exist",
     );
   });

test/abi/airdrop.abi.json

@@ -1,7 +1,6 @@
 [
   {
     "inputs": [],
-    "payable": false,
     "stateMutability": "nonpayable",
     "type": "constructor"
   },
@@ -136,7 +135,6 @@
     "type": "event"
   },
   {
-    "constant": true,
     "inputs": [
       {
         "internalType": "uint256",
@@ -157,12 +155,10 @@
         "type": "uint256"
       }
     ],
-    "payable": false,
     "stateMutability": "view",
     "type": "function"
   },
   {
-    "constant": false,
     "inputs": [
       {
         "internalType": "uint256",
@@ -183,12 +179,10 @@
         "type": "bool"
       }
     ],
-    "payable": false,
     "stateMutability": "nonpayable",
     "type": "function"
   },
   {
-    "constant": false,
     "inputs": [
       {
         "internalType": "uint256",
@@ -204,12 +198,10 @@
         "type": "bool"
       }
     ],
-    "payable": false,
     "stateMutability": "nonpayable",
     "type": "function"
   },
   {
-    "constant": false,
     "inputs": [
       {
         "internalType": "uint256",
@@ -235,12 +227,10 @@
         "type": "bool"
       }
     ],
-    "payable": false,
     "stateMutability": "nonpayable",
     "type": "function"
   },
   {
-    "constant": true,
     "inputs": [
       {
         "internalType": "uint256",
@@ -256,12 +246,10 @@
         "type": "uint256"
       }
     ],
-    "payable": false,
     "stateMutability": "view",
     "type": "function"
   },
   {
-    "constant": true,
     "inputs": [
       {
         "internalType": "uint256",
@@ -302,12 +290,10 @@
         "type": "uint256"
       }
     ],
-    "payable": false,
     "stateMutability": "view",
     "type": "function"
   },
   {
-    "constant": true,
     "inputs": [],
     "name": "nextStreamId",
     "outputs": [
@@ -317,12 +303,10 @@
         "type": "uint256"
       }
     ],
-    "payable": false,
     "stateMutability": "view",
     "type": "function"
   },
   {
-    "constant": true,
     "inputs": [],
     "name": "torn",
     "outputs": [
@@ -332,12 +316,10 @@
         "type": "address"
       }
     ],
-    "payable": false,
     "stateMutability": "view",
     "type": "function"
   },
   {
-    "constant": true,
     "inputs": [],
     "name": "tornadoGovernance",
     "outputs": [
@@ -347,22 +329,15 @@
         "type": "address"
       }
     ],
-    "payable": false,
     "stateMutability": "view",
     "type": "function"
   },
   {
-    "constant": false,
     "inputs": [
       {
         "internalType": "uint256",
         "name": "streamId",
         "type": "uint256"
-      },
-      {
-        "internalType": "uint256",
-        "name": "amount",
-        "type": "uint256"
       }
     ],
     "name": "withdrawFromStream",
@@ -373,12 +348,10 @@
         "type": "bool"
       }
     ],
-    "payable": false,
     "stateMutability": "nonpayable",
     "type": "function"
   },
   {
-    "constant": false,
     "inputs": [],
     "name": "withdrawFunds",
     "outputs": [
@@ -388,7 +361,6 @@
         "type": "bool"
       }
     ],
-    "payable": false,
     "stateMutability": "nonpayable",
     "type": "function"
   }

test/utils.js

@@ -2,8 +2,6 @@ const { ethers, network } = require("hardhat");
 const { time } = require("@nomicfoundation/hardhat-toolbox/network-helpers");
 
 const stakingAddr = "0x5B3f656C80E8ddb9ec01Dd9018815576E9238c29";
-const gasCompensationAddr = "0xFA4C1f3f7D5dd7c12a9Adb82Cd7dDA542E3d59ef";
-const userVaultAddr = "0x2F50508a8a3D323B91336FA3eA6ae50E55f32185";
 const governanceAddr = "0x5efda50f22d34F262c29268506C5Fa42cB56A1Ce";
 const tornAddr = "0x77777FeDdddFfC19Ff86DB637967013e6C6A116C";