From bcdabd0499b05f26217e5178dbe3be9fb4c1606f Mon Sep 17 00:00:00 2001
From: Micky Mousse <newmickymousse@pm.me>
Date: Fri, 4 Mar 2022 19:18:16 -0500
Subject: [PATCH 1/5] feat: multirewards ssbeets

---
 contracts/Strategy.sol  | 149 +++++++++++++++++++++-------------------
 tests/conftest.py       |  41 +++++++----
 tests/test_fix.py       |  78 ---------------------
 tests/test_migration.py |  67 +-----------------
 tests/test_operation.py |   9 ++-
 tests/util.py           |  14 +++-
 6 files changed, 128 insertions(+), 230 deletions(-)
 delete mode 100644 tests/test_fix.py

diff --git a/contracts/Strategy.sol b/contracts/Strategy.sol
index 4a2cfbf..0cee483 100644
--- a/contracts/Strategy.sol
+++ b/contracts/Strategy.sol
@@ -17,9 +17,10 @@ contract Strategy is BaseStrategy {
 
     IBalancerVault public balancerVault;
     IBalancerPool public bpt;
-    IERC20 public rewardToken;
+
+    IERC20[] public rewardTokens;
     IAsset[] internal assets;
-    SwapSteps internal swapSteps;
+    SwapSteps[] internal swapSteps;
     bytes32 public balancerPoolId;
     uint8 internal numTokens;
     uint8 internal tokenIndex;
@@ -39,12 +40,13 @@ contract Strategy is BaseStrategy {
     }
 
     struct Toggles {
-        bool claimRewards;
+        bool doClaimRewards;
         bool doSellRewards;
         bool abandonRewards;
     }
 
     uint256 internal constant max = type(uint256).max;
+    IERC20 private constant beets = IERC20(0xF24Bcf4d1e507740041C9cFd2DddB29585aDCe1e);
 
     //1	    0.01%
     //5	    0.05%
@@ -155,7 +157,7 @@ contract Strategy is BaseStrategy {
         keep = governance();
         keepBips = 1000;
 
-        toggles = Toggles({claimRewards : true, doSellRewards : true, abandonRewards : false});
+        toggles = Toggles({doClaimRewards : true, doSellRewards : true, abandonRewards : false});
     }
 
     // ******** OVERRIDE THESE METHODS FROM BASE CONTRACT ************
@@ -169,7 +171,7 @@ contract Strategy is BaseStrategy {
     }
 
     function prepareReturn(uint256 _debtOutstanding) internal override returns (uint256 _profit, uint256 _loss, uint256 _debtPayment){
-        if(toggles.claimRewards){
+        if(toggles.doClaimRewards){
             _claimRewards();
         }
         if (toggles.doSellRewards) {
@@ -213,17 +215,15 @@ contract Strategy is BaseStrategy {
             return;
         }
 
-        // put want into lp then put want-lp into masterchef
-        if (_joinPool()) {
+        // put want into lp then put loose want-lp into masterchef
+        _joinPool();
+        uint256 _balanceOfBpt = balanceOfBpt();
+        if (_balanceOfBpt > 0) {
             // put all want-lp into masterchef
-            _depositIntoMasterChef(balanceOfBpt());
+            _depositIntoMasterChef(_balanceOfBpt);
         }
     }
 
-    function depositIntoMasterChef(uint _bpts) external onlyVaultManagers {
-        _depositIntoMasterChef(_bpts);
-    }
-
     function _depositIntoMasterChef(uint _bpts) internal {
         masterChef.deposit(masterChefPoolId, _bpts, address(this));
     }
@@ -235,7 +235,6 @@ contract Strategy is BaseStrategy {
             uint256 toExitAmount = tokensToBpts(_amountNeeded.sub(looseAmount));
             // withdraw needed bpt out of masterchef and sell it for want
             _withdrawFromMasterChefAndSellBpt(toExitAmount);
-
             _liquidatedAmount = Math.min(balanceOfWant(), _amountNeeded);
             _loss = _amountNeeded.sub(_liquidatedAmount);
         } else {
@@ -258,9 +257,12 @@ contract Strategy is BaseStrategy {
         if (_balanceOfBpt > 0) {
             bpt.transfer(_newStrategy, _balanceOfBpt);
         }
-        uint256 rewards = balanceOfReward();
-        if (rewards > 0) {
-            rewardToken.safeTransfer(_newStrategy, rewards);
+        for (uint i = 0; i < rewardTokens.length; i++) {
+            IERC20 token = rewardTokens[i];
+            uint256 balance = token.balanceOf(address(this));
+            if (balance > 0) {
+                token.safeTransfer(_newStrategy, balance);
+            }
         }
     }
 
@@ -300,50 +302,51 @@ contract Strategy is BaseStrategy {
         _sellBpt(_amountBpt);
     }
 
-    function claimRewards() external onlyVaultManagers {
-        _claimRewards();
-    }
-
     // claim all beets rewards from masterchef
     function _claimRewards() internal {
         if (getPendingBeets() > 0) {
-            uint256 rewardBal = balanceOfReward();
+            uint256 prevBeetsBal = balanceOfBeets();
             masterChef.harvest(masterChefPoolId, address(this));
-            uint256 keepBal = balanceOfReward().sub(rewardBal).mul(keepBips).div(basisOne);
+            uint256 keepBal = balanceOfBeets().sub(prevBeetsBal).mul(keepBips).div(basisOne);
             if (keepBal > 0) {
-                rewardToken.safeTransfer(keep, keepBal);
+                beets.safeTransfer(keep, keepBal);
             }
         }
     }
 
-    function sellRewards() external onlyVaultManagers {
-        _sellRewards();
-    }
-
     function _sellRewards() internal {
-        uint256 amount = balanceOfReward();
-        if (amount > 0) {
-            uint length = swapSteps.poolIds.length;
-            IBalancerVault.BatchSwapStep[] memory steps = new IBalancerVault.BatchSwapStep[](length);
-            int[] memory limits = new int[](length + 1);
-            limits[0] = int(amount);
-            for (uint j = 0; j < length; j++) {
-                steps[j] = IBalancerVault.BatchSwapStep(swapSteps.poolIds[j],
-                    j,
-                    j + 1,
-                    j == 0 ? amount : 0,
-                    abi.encode(0)
-                );
+        for (uint8 i = 0; i < rewardTokens.length; i++) {
+            ERC20 rewardToken = ERC20(address(rewardTokens[i]));
+            uint256 amount = rewardToken.balanceOf(address(this));
+
+            uint decReward = rewardToken.decimals();
+            uint decWant = ERC20(address(want)).decimals();
+
+            if (amount > 10 ** (decReward > decWant ? decReward.sub(decWant) : 0)) {
+                uint length = swapSteps[i].poolIds.length;
+                IBalancerVault.BatchSwapStep[] memory steps = new IBalancerVault.BatchSwapStep[](length);
+                int[] memory limits = new int[](length + 1);
+                limits[0] = int(amount);
+                for (uint j = 0; j < length; j++) {
+                    steps[j] = IBalancerVault.BatchSwapStep(swapSteps[i].poolIds[j],
+                        j,
+                        j + 1,
+                        j == 0 ? amount : 0,
+                        abi.encode(0)
+                    );
+                }
+                balancerVault.batchSwap(IBalancerVault.SwapKind.GIVEN_IN,
+                    steps,
+                    swapSteps[i].assets,
+                    IBalancerVault.FundManagement(address(this), false, address(this), false),
+                    limits,
+                    now + 10);
             }
-            balancerVault.batchSwap(IBalancerVault.SwapKind.GIVEN_IN,
-                steps,
-                swapSteps.assets,
-                IBalancerVault.FundManagement(address(this), false, address(this), false),
-                limits,
-                now + 10);
         }
     }
 
+
+
     function balanceOfWant() public view returns (uint256 _amount){
         return want.balanceOf(address(this));
     }
@@ -360,8 +363,8 @@ contract Strategy is BaseStrategy {
         return masterChef.pendingBeets(masterChefPoolId, address(this));
     }
 
-    function balanceOfReward() public view returns (uint256 _amount){
-        return rewardToken.balanceOf(address(this));
+    function balanceOfBeets() public view returns (uint256 _amount){
+        return beets.balanceOf(address(this));
     }
 
     // returns an estimate of want tokens based on bpt balance
@@ -370,12 +373,12 @@ contract Strategy is BaseStrategy {
     }
 
     /// use bpt rate to estimate equivalent amount of want.
-    function bptsToTokens(uint _amountBpt) public view returns (uint _amount){
+    function bptsToTokens(uint _amountBpt) internal view returns (uint _amount){
         uint unscaled = _amountBpt.mul(bpt.getRate()).div(1e18);
         return _scaleDecimals(unscaled, ERC20(address(bpt)), ERC20(address(want)));
     }
 
-    function tokensToBpts(uint _amountTokens) public view returns (uint _amount){
+    function tokensToBpts(uint _amountTokens) internal view returns (uint _amount){
         uint unscaled = _amountTokens.mul(1e18).div(bpt.getRate());
         return _scaleDecimals(unscaled, ERC20(address(want)), ERC20(address(bpt)));
     }
@@ -399,10 +402,6 @@ contract Strategy is BaseStrategy {
         );
     }
 
-    // this allows us to also sell bpt externally
-    function sellBpt(uint256 _amountBpts) external onlyVaultManagers {
-        _sellBpt(_amountBpts);
-    }
 
     // sell bpt for want at current bpt rate
     function _sellBpt(uint256 _amountBpts) internal {
@@ -432,14 +431,25 @@ contract Strategy is BaseStrategy {
         return false;
     }
 
-    function whitelistReward(address _rewardToken, SwapSteps memory _steps) public onlyVaultManagers {
-        rewardToken = IERC20(_rewardToken);
-        rewardToken.approve(address(balancerVault), 0);
-        rewardToken.approve(address(balancerVault), max);
-        swapSteps = _steps;
+    // for partnership rewards like TUSD or airdrops
+    function whitelistRewards(address _rewardToken, SwapSteps memory _steps) public onlyVaultManagers {
+        IERC20 token = IERC20(_rewardToken);
+        token.approve(address(balancerVault), max);
+        rewardTokens.push(token);
+        swapSteps.push(_steps);
+    }
+
+    function delistAllRewards() public onlyVaultManagers {
+        for (uint i = 0; i < rewardTokens.length; i++) {
+            rewardTokens[i].approve(address(balancerVault), 0);
+        }
+        IERC20[] memory noRewardTokens;
+        rewardTokens = noRewardTokens;
+        delete swapSteps;
     }
 
-    function setParams(uint256 _maxSlippageIn, uint256 _maxSlippageOut, uint256 _maxSingleDeposit, uint256 _minDepositPeriod) public onlyVaultManagers {
+    function setParams(uint256 _maxSlippageIn, uint256 _maxSlippageOut, uint256 _maxSingleDeposit, uint256 _minDepositPeriod,
+            address _keep, uint _keepBips) public onlyVaultManagers {
         require(_maxSlippageIn <= basisOne);
         maxSlippageIn = _maxSlippageIn;
 
@@ -448,24 +458,23 @@ contract Strategy is BaseStrategy {
 
         maxSingleDeposit = _maxSingleDeposit;
         minDepositPeriod = _minDepositPeriod;
+
+        require(_keepBips <= basisOne);
+        keep = _keep;
+        keepBips = _keepBips;
     }
 
-    function setToggles(bool _claimRewards, bool _doSellRewards, bool _abandon) external onlyVaultManagers {
-        toggles.claimRewards = _claimRewards;
+    function setToggles(bool _doClaimRewards, bool _doSellRewards, bool _abandon) external onlyVaultManagers {
+        toggles.doClaimRewards = _doClaimRewards;
         toggles.doSellRewards = _doSellRewards;
         toggles.abandonRewards = _abandon;
     }
 
     // swap step contains information on multihop sells
-    function getSwapSteps() public view returns (SwapSteps memory){
-        return swapSteps;
-    }
+    function getSwapSteps() public view returns (SwapSteps[] memory){
+         return swapSteps;
+     }
 
-    function setKeepParams(address _keep, uint _keepBips) external onlyGovernance {
-        require(keepBips <= basisOne);
-        keep = _keep;
-        keepBips = _keepBips;
-    }
 
     // Balancer requires this contract to be payable, so we add ability to sweep stuck ETH
     function sweepETH() public onlyGovernance {
diff --git a/tests/conftest.py b/tests/conftest.py
index 0a94926..c556745 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -52,6 +52,7 @@ def token():
     # 0x049d68029688eAbF473097a2fC38ef61633A3C7A fUSDT
     # 0x04068DA6C83AFCFA0e13ba15A6696662335D5B75 USDC
     # 0x82f0B8B456c1A451378467398982d4834b6829c1 MIM
+    # 0x9879aBDea01a879644185341F7aF7d8343556B7a TUSD
     token_address = "0x04068DA6C83AFCFA0e13ba15A6696662335D5B75"
     yield Contract(token_address)
 
@@ -62,7 +63,8 @@ def token2():
     # 0x049d68029688eAbF473097a2fC38ef61633A3C7A fUSDT
     # 0x04068DA6C83AFCFA0e13ba15A6696662335D5B75 USDC
     # 0x82f0B8B456c1A451378467398982d4834b6829c1 MIM
-    token_address = "0x8D11eC38a3EB5E956B052f67Da8Bdc9bef8Abf3E"
+    # 0x9879aBDea01a879644185341F7aF7d8343556B7a TUSD
+    token_address = "0x9879aBDea01a879644185341F7aF7d8343556B7a"
     yield Contract(token_address)
 
 
@@ -72,6 +74,7 @@ def token_whale(accounts):
     # 0x2dd7C9371965472E5A5fD28fbE165007c61439E1 fUSDT
     # 0x93C08a3168fC469F3fC165cd3A471D19a37ca19e USDC
     # 0x8D9AED9882b4953a0c9fa920168fa1FDfA0eBE75 DAI
+    # 0x789B5DBd47d7Ca3799f8E9FdcE01bC5E356fcDF1 TUSD
     return accounts.at("0x93C08a3168fC469F3fC165cd3A471D19a37ca19e", force=True)
 
 
@@ -81,12 +84,13 @@ def token2_whale(accounts):
     # 0x2dd7C9371965472E5A5fD28fbE165007c61439E1 fUSDT
     # 0x93C08a3168fC469F3fC165cd3A471D19a37ca19e USDC
     # 0x8D9AED9882b4953a0c9fa920168fa1FDfA0eBE75 DAI
-    return accounts.at("0x8D9AED9882b4953a0c9fa920168fa1FDfA0eBE75", force=True)
+    # 0x789B5DBd47d7Ca3799f8E9FdcE01bC5E356fcDF1 TUSD
+    return accounts.at("0x789B5DBd47d7Ca3799f8E9FdcE01bC5E356fcDF1", force=True)
 
 
 @pytest.fixture
 def amount(accounts, token, user, token_whale):
-    amount = 1_000_000 * 10 ** token.decimals()
+    amount = 100_000 * 10 ** token.decimals()
     # In order to get some funds for the token you are about to use,
     token.transfer(user, amount, {"from": token_whale})
     yield amount
@@ -105,6 +109,10 @@ def wftm():
     token_address = "0x21be370D5312f44cB42ce377BC9b8a0cEF1A4C83"
     yield Contract(token_address)
 
+@pytest.fixture
+def tusd():
+    token_address = "0x9879aBDea01a879644185341F7aF7d8343556B7a"
+    yield Contract(token_address)
 
 @pytest.fixture
 def beets():
@@ -150,7 +158,8 @@ def balancer_vault():
 def pool():
     # 0xD163415BD34EF06f57C58D2AEd5A5478AfB464cC MIM-USDC-USDT Stable Pool
     # 0xeCAa1cBd28459d34B766F9195413Cb20122Fb942 dai-usdc stable pool
-    address = "0xeCAa1cBd28459d34B766F9195413Cb20122Fb942"
+    # 0xcf9d4940fe4c194c83d4d3b1de4c2dff4233f612 tusd-usdc stable pool
+    address = "0xcf9d4940fe4c194c83d4d3b1de4c2dff4233f612"
     yield Contract(address)
 
 
@@ -171,16 +180,23 @@ def beetsUsdcPoolId():
     yield 0x03c6b3f09d2504606936b1a4decefad204687890000200000000000000000015
 
 
+# @pytest.fixture
+# def usdcTokenPoolId():
+#     id = 0xecaa1cbd28459d34b766f9195413cb20122fb942000200000000000000000120  # usdc-mim
+#     yield id
+
 @pytest.fixture
-def usdcTokenPoolId():
-    id = 0xecaa1cbd28459d34b766f9195413cb20122fb942000200000000000000000120  # usdc-mim
+def tusdTokenPoolId():
+    id = 0xcf9d4940fe4c194c83d4d3b1de4c2dff4233f612000200000000000000000253  # usdc-tusd
     yield id
 
-
 @pytest.fixture
 def swapStepsBeets(beetsUsdcPoolId, beets, token):
     yield ([beetsUsdcPoolId], [beets, token])
 
+@pytest.fixture
+def swapStepsTusd(tusdTokenPoolId, tusd, token):
+    yield ([tusdTokenPoolId], [tusd, token])
 
 @pytest.fixture
 def strategyFactory(strategist, keeper, vault, StrategyFactory, gov, balancer_vault, pool, beets, usdc, beetsUsdcPool,
@@ -188,16 +204,17 @@ def strategyFactory(strategist, keeper, vault, StrategyFactory, gov, balancer_va
                     masterChef,
                     swapStepsBeets):
     factory = strategist.deploy(StrategyFactory, vault, balancer_vault, pool, masterChef, 5, 5, 100_000, 2 * 60 * 60,
-                                33)
+                                58)
     yield factory
 
 
 @pytest.fixture
-def strategy(strategist, keeper, vault, strategyFactory, gov, balancer_vault, pool, beets, usdc, beetsUsdcPool,
-             management, masterChef, swapStepsBeets, Strategy):
+def strategy(strategist, keeper, vault, strategyFactory, gov, balancer_vault, pool, beets, usdc,
+             management, masterChef, swapStepsBeets, tusd, swapStepsTusd, Strategy):
     strategy = Strategy.at(strategyFactory.original())
     strategy.setKeeper(keeper,{'from': gov})
-    strategy.whitelistReward(beets, swapStepsBeets, {'from': gov})
+    strategy.whitelistRewards(beets, swapStepsBeets, {'from': gov})
+    strategy.whitelistRewards(tusd, swapStepsTusd, {'from': gov})
     vault.addStrategy(strategy, 10_000, 0, 2 ** 256 - 1, 1_000, {"from": gov})
     vault.setManagementFee(0, {"from": gov})
     chain.sleep(1)
@@ -208,4 +225,4 @@ def strategy(strategist, keeper, vault, strategyFactory, gov, balancer_vault, po
 def RELATIVE_APPROX():
     # making this more lenient bc of single
     # sided deposits incurring slippage
-    yield 1e-3
\ No newline at end of file
+    yield 1e-3
diff --git a/tests/test_fix.py b/tests/test_fix.py
deleted file mode 100644
index a440a35..0000000
--- a/tests/test_fix.py
+++ /dev/null
@@ -1,78 +0,0 @@
-import pytest
-from brownie import accounts, Contract
-import util
-
-def test_real_migration_and_multiple_harvest(
-        chain,
-        vault,
-        amount,
-        Strategy,
-        strategist,
-        gov,
-        user,
-        masterChef,
-        RELATIVE_APPROX,
-        balancer_vault, pool, management, swapStepsBeets, beets
-):
-    old = Contract("0xF864f92e88054AA05639324090b411C1D55B4a5B")
-    token = Contract(old.want())
-    vault = Contract(old.vault())
-    gov = accounts.at(vault.governance(), force=True)
-    fromGov = {'from': gov}
-
-    util.stateOfStrat("old strategy before migration", old, token)
-
-    fixed_strategy = strategist.deploy(Strategy, vault, balancer_vault, pool, masterChef, 5, 5, 100_000, 2 * 60 * 60,
-                                       33)
-    fixed_strategy.whitelistReward(beets, swapStepsBeets, fromGov)
-
-
-
-    # steady beets 2 pool = #33
-    old_pending = masterChef.pendingBeets(33, old)
-    print(f'pending beets from old strategy: {old_pending}')
-
-    vault.migrateStrategy(old, fixed_strategy, fromGov)
-    util.stateOfStrat("old strategy after migration", old, token)
-    util.stateOfStrat("new strategy after migration", fixed_strategy, token)
-
-    total_debt = vault.strategies(fixed_strategy)["totalDebt"]
-    assert fixed_strategy.estimatedTotalAssets() >= total_debt
-
-    # this lets gov sweep and take x% cut
-    assert beets.balanceOf(fixed_strategy) >= old_pending
-    print(f'vault state: {vault.strategies(fixed_strategy)}')
-
-    fixed_strategy.depositIntoMasterChef(fixed_strategy.balanceOfBpt(), fromGov)
-    print(f'vault state: {vault.strategies(fixed_strategy)}')
-
-    util.stateOfStrat("old strategy after harvest", old, token)
-    util.stateOfStrat("new strategy after harvest", fixed_strategy, token)
-
-    tx = fixed_strategy.harvest(fromGov)
-    tx = fixed_strategy.harvest(fromGov)
-
-    # sell everything
-    fixed_strategy.setEmergencyExit(fromGov)
-    # fixed_strategy.updateStrategyDebtRatio(0,fromGov)
-    # rewards sold separately
-    # fixed_strategy.sellRewards(fromGov)
-
-    fixed_strategy.setParams(
-        2_000,
-        2_000,
-        fixed_strategy.maxSingleDeposit(),
-        fixed_strategy.minDepositPeriod(),
-        fromGov
-    )
-    fixed_strategy.setDoHealthCheck(False, fromGov)
-    chain.sleep(1)
-    fixed_strategy.harvest(fromGov)
-
-    print(f'vault state: {vault.strategies(fixed_strategy)}')
-
-    util.stateOfStrat("debt ratio 0", fixed_strategy, token)
-
-    # hopefully the gains from trading fees cancels out slippage
-    print(f'net loss from exit: {vault.strategies(fixed_strategy)["totalLoss"]}')
-    assert fixed_strategy.estimatedTotalAssets() == 0
diff --git a/tests/test_migration.py b/tests/test_migration.py
index 749cc40..8dd778b 100644
--- a/tests/test_migration.py
+++ b/tests/test_migration.py
@@ -28,71 +28,6 @@ def test_migration(
 
     # migrate to a new strategy
     new_strategy = strategist.deploy(Strategy, vault, balancer_vault, pool, masterChef, 5, 5, 100_000, 2 * 60 * 60,
-                                     33)
+                                     58)
     vault.migrateStrategy(strategy, new_strategy, {"from": gov})
     assert (pytest.approx(new_strategy.estimatedTotalAssets(), rel=RELATIVE_APPROX) == amount)
-
-
-def test_real_migration(
-        chain,
-        vault,
-        amount,
-        Strategy,
-        strategist,
-        gov,
-        user,
-        masterChef,
-        RELATIVE_APPROX,
-        balancer_vault, pool, management, swapStepsBeets, beets
-):
-    old = Contract("0xF864f92e88054AA05639324090b411C1D55B4a5B")
-    token = Contract(old.want())
-    vault = Contract(old.vault())
-    gov = accounts.at(vault.governance(), force=True)
-    fromGov = {'from': gov}
-
-    util.stateOfStrat("old strategy before migration", old, token)
-
-    fixed_strategy = strategist.deploy(Strategy, vault, balancer_vault, pool, masterChef, 5, 5, 100_000, 2 * 60 * 60,
-                                       33)
-    fixed_strategy.whitelistReward(beets, swapStepsBeets, fromGov)
-
-    # steady beets 2 pool = #33
-    old_pending = masterChef.pendingBeets(33, old)
-    print(f'pending beets from old strategy: {old_pending}')
-
-    vault.migrateStrategy(old, fixed_strategy, fromGov)
-    fixed_strategy.depositIntoMasterChef(fixed_strategy.balanceOfBptInMasterChef(), fromGov)
-    util.stateOfStrat("old strategy after migration", old, token)
-    util.stateOfStrat("new strategy after migration", fixed_strategy, token)
-    total_debt = vault.strategies(fixed_strategy)["totalDebt"]
-    assert fixed_strategy.estimatedTotalAssets() >= total_debt
-
-    # this lets gov sweep and take x% cut
-    assert beets.balanceOf(fixed_strategy) >= old_pending
-    print(f'vault state: {vault.strategies(fixed_strategy)}')
-
-    fixed_strategy.depositIntoMasterChef(fixed_strategy.balanceOfBpt(), fromGov)
-    print(f'vault state: {vault.strategies(fixed_strategy)}')
-
-    util.stateOfStrat("old strategy after harvest", old, token)
-    util.stateOfStrat("new strategy after harvest", fixed_strategy, token)
-
-    # sell everything
-    fixed_strategy.setEmergencyExit(fromGov)
-    # rewards sold separately
-    fixed_strategy.sellRewards(fromGov)
-    with brownie.reverts():
-        fixed_strategy.harvest(fromGov)
-
-    fixed_strategy.setParams(fixed_strategy.maxSlippageIn(), 10000, fixed_strategy.maxSingleDeposit(), fixed_strategy.minDepositPeriod(), fromGov)
-    chain.sleep(1)
-    fixed_strategy.harvest(fromGov)
-
-    print(f'vault state: {vault.strategies(fixed_strategy)}')
-
-    util.stateOfStrat("debt ratio 0", fixed_strategy, token)
-
-    # hopefully the gains from trading fees cancels out slippage
-    print(f'net loss from exit: {vault.strategies(fixed_strategy)["totalLoss"]}')
-    assert fixed_strategy.estimatedTotalAssets() == 0
diff --git a/tests/test_operation.py b/tests/test_operation.py
index 22b2f34..3d7ac72 100644
--- a/tests/test_operation.py
+++ b/tests/test_operation.py
@@ -84,6 +84,8 @@ def test_profitable_harvest(
 
     before_pps = vault.pricePerShare()
 
+    chain.mine(1)
+    chain.sleep(1)
     # Harvest 2: Realize profit
     util.airdrop_rewards(strategy, beets, beets_whale)
 
@@ -116,7 +118,7 @@ def test_deposit_all(chain, token, vault, strategy, user, strategist, amount, RE
         strategy.tend({'from': gov})
         util.stateOfStrat("tend", strategy, token)
         assert pytest.approx(strategy.estimatedTotalAssets(), rel=RELATIVE_APPROX) == amount
-        chain.sleep(strategy.minDepositPeriod() + 1)
+        # chain.sleep(strategy.minDepositPeriod() + 1)
         chain.mine(1)
 
     before_pps = vault.pricePerShare()
@@ -138,6 +140,9 @@ def test_deposit_all(chain, token, vault, strategy, user, strategist, amount, RE
 
     vault.updateStrategyDebtRatio(strategy.address, 5_000, {"from": gov})
     chain.sleep(1)
+
+    # strategy.setDoHealthCheck(False, {'from':gov})
+
     strategy.harvest({"from": strategist})
     util.stateOfStrat("after harvest 5000", strategy, token)
 
@@ -328,7 +333,7 @@ def test_unbalanced_pool_withdraw(chain, token, vault, strategy, user, strategis
     old_slippage = strategy.maxSlippageOut()
 
     # loosen the slippage check to let the lossy withdraw go through
-    strategy.setParams(10000, 10000, strategy.maxSingleDeposit(), strategy.minDepositPeriod(), {'from': gov})
+    strategy.setParams(10000, 10000, strategy.maxSingleDeposit(), strategy.minDepositPeriod(), strategy.keep(), strategy.keepBips(), {'from': gov})
     vault.withdraw(vault.balanceOf(user) / 2, user, 10000, {"from": user})
     print(f'pool state: {balancer_vault.getPoolTokens(pool.getPoolId())}')
     print(f'user balance: {token.balanceOf(user)}')
diff --git a/tests/util.py b/tests/util.py
index 983c481..cac9efa 100644
--- a/tests/util.py
+++ b/tests/util.py
@@ -3,7 +3,17 @@
 
 def airdrop_rewards(strategy, beets, beets_whale):
     beets.approve(strategy, 2 ** 256 - 1, {'from': beets_whale})
-    beets.transfer(strategy, 3_000 * 1e18, {'from': beets_whale})
+    beets.transfer(strategy, 500 * 1e18, {'from': beets_whale})
+
+def airdrop_tusd_rewards(strategy, token2, token2_whale):
+    token2.approve(strategy, 2 ** 256 - 1, {'from': token2_whale})
+    token2.transfer(strategy, 3_000 * 1e18, {'from': token2_whale})
+
+def airdrop_all_rewards(strategy, beets, beets_whale, token2, token2_whale):
+    token2.approve(strategy, 2 ** 256 - 1, {'from': token2_whale})
+    token2.transfer(strategy, 3_000 * 1e18, {'from': token2_whale})
+    token2.approve(strategy, 2 ** 256 - 1, {'from': token2_whale})
+    token2.transfer(strategy, 3_000 * 1e18, {'from': token2_whale})
 
 
 def stateOfStrat(msg, strategy, token):
@@ -12,5 +22,5 @@ def stateOfStrat(msg, strategy, token):
     print(f'Balance of {token.symbol()}: {strategy.balanceOfWant() / wantDec}')
     print(f'Balance of Bpt: {strategy.balanceOfBpt() / 1e18}')
     print(f'balanceOfBptInMasterChef: {strategy.balanceOfBptInMasterChef() / 1e18}')
-    print(f'Balance of BEETS: {strategy.balanceOfReward()/ wantDec}')
+    print(f'Balance of BEETS: {strategy.balanceOfBeets()/ wantDec}')
     print(f'Estimated Total Assets: {strategy.estimatedTotalAssets() / wantDec}')

From 69a538bb6d83899f0d61f0b8b79704c95b7462c1 Mon Sep 17 00:00:00 2001
From: Micky Mousse <newmickymousse@pm.me>
Date: Sat, 5 Mar 2022 13:12:13 -0500
Subject: [PATCH 2/5] fix: reviewers comments

---
 contracts/Strategy.sol | 27 ++++++++++++++++++++++-----
 tests/test_revoke.py   |  2 ++
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/contracts/Strategy.sol b/contracts/Strategy.sol
index 0cee483..fab6405 100644
--- a/contracts/Strategy.sol
+++ b/contracts/Strategy.sol
@@ -34,6 +34,15 @@ contract Strategy is BaseStrategy {
     // masterchef
     IBeethovenxMasterChef public masterChef;
 
+    modifier isVaultManager {
+        checkVaultManagers();
+        _;
+    }
+
+    function checkVaultManagers() internal {
+        require(msg.sender == vault.governance() || msg.sender == vault.management());
+    }
+
     struct SwapSteps {
         bytes32[] poolIds;
         IAsset[] assets;
@@ -276,7 +285,7 @@ contract Strategy is BaseStrategy {
 
     // HELPERS //
 
-    function withdrawFromMasterChef(uint256 _amountBpt) external onlyVaultManagers {
+    function withdrawFromMasterChef(uint256 _amountBpt) external isVaultManager {
         _withdrawFromMasterChef(address(this), _amountBpt);
     }
 
@@ -302,6 +311,10 @@ contract Strategy is BaseStrategy {
         _sellBpt(_amountBpt);
     }
 
+    function claimRewards() external isVaultManager {
+        _claimRewards();
+    }
+
     // claim all beets rewards from masterchef
     function _claimRewards() internal {
         if (getPendingBeets() > 0) {
@@ -314,6 +327,10 @@ contract Strategy is BaseStrategy {
         }
     }
 
+    function sellRewards() external isVaultManager {
+        _sellRewards();
+    }
+
     function _sellRewards() internal {
         for (uint8 i = 0; i < rewardTokens.length; i++) {
             ERC20 rewardToken = ERC20(address(rewardTokens[i]));
@@ -432,14 +449,14 @@ contract Strategy is BaseStrategy {
     }
 
     // for partnership rewards like TUSD or airdrops
-    function whitelistRewards(address _rewardToken, SwapSteps memory _steps) public onlyVaultManagers {
+    function whitelistRewards(address _rewardToken, SwapSteps memory _steps) public isVaultManager {
         IERC20 token = IERC20(_rewardToken);
         token.approve(address(balancerVault), max);
         rewardTokens.push(token);
         swapSteps.push(_steps);
     }
 
-    function delistAllRewards() public onlyVaultManagers {
+    function delistAllRewards() public isVaultManager {
         for (uint i = 0; i < rewardTokens.length; i++) {
             rewardTokens[i].approve(address(balancerVault), 0);
         }
@@ -449,7 +466,7 @@ contract Strategy is BaseStrategy {
     }
 
     function setParams(uint256 _maxSlippageIn, uint256 _maxSlippageOut, uint256 _maxSingleDeposit, uint256 _minDepositPeriod,
-            address _keep, uint _keepBips) public onlyVaultManagers {
+            address _keep, uint _keepBips) public isVaultManager {
         require(_maxSlippageIn <= basisOne);
         maxSlippageIn = _maxSlippageIn;
 
@@ -464,7 +481,7 @@ contract Strategy is BaseStrategy {
         keepBips = _keepBips;
     }
 
-    function setToggles(bool _doClaimRewards, bool _doSellRewards, bool _abandon) external onlyVaultManagers {
+    function setToggles(bool _doClaimRewards, bool _doSellRewards, bool _abandon) external isVaultManager {
         toggles.doClaimRewards = _doClaimRewards;
         toggles.doSellRewards = _doSellRewards;
         toggles.abandonRewards = _abandon;
diff --git a/tests/test_revoke.py b/tests/test_revoke.py
index bed34ff..09255f2 100644
--- a/tests/test_revoke.py
+++ b/tests/test_revoke.py
@@ -14,6 +14,7 @@ def test_revoke_strategy_from_vault(
     # In order to pass this tests, you will need to implement prepareReturn.
     vault.revokeStrategy(strategy.address, {"from": gov})
     chain.sleep(1)
+    strategy.setDoHealthCheck(False, {"from": gov})
     strategy.harvest({"from": gov})
     assert pytest.approx(token.balanceOf(vault.address), rel=RELATIVE_APPROX) == amount
 
@@ -30,5 +31,6 @@ def test_revoke_strategy_from_strategy(
 
     strategy.setEmergencyExit({"from": gov})
     chain.sleep(1)
+    strategy.setDoHealthCheck(False, {"from": gov})
     strategy.harvest({"from": gov})
     assert pytest.approx(token.balanceOf(vault.address), rel=RELATIVE_APPROX) == amount

From a168d09a11aba812aedca33ccb26ae0fb6252903 Mon Sep 17 00:00:00 2001
From: Micky Mousse <newmickymousse@pm.me>
Date: Tue, 15 Mar 2022 14:35:58 -0400
Subject: [PATCH 3/5] fix: peer reviewers comments

---
 contracts/Strategy.sol  | 15 ++++---
 tests/conftest.py       |  5 ++-
 tests/test_operation.py | 95 ++++++++++++++++++++++++++++++-----------
 tests/util.py           | 16 +++----
 4 files changed, 90 insertions(+), 41 deletions(-)

diff --git a/contracts/Strategy.sol b/contracts/Strategy.sol
index fab6405..5afb0a5 100644
--- a/contracts/Strategy.sol
+++ b/contracts/Strategy.sol
@@ -7,13 +7,19 @@ pragma experimental ABIEncoderV2;
 
 // These are the core Yearn libraries
 import {BaseStrategy, StrategyParams} from "@yearnvaults/contracts/BaseStrategy.sol";
-import {Address} from "@openzeppelin/contracts/token/ERC20/SafeERC20.sol";
-import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
-import {Math} from "@openzeppelin/contracts/math/Math.sol";
+import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
+import "@openzeppelin/contracts/math/SafeMath.sol";
+import "@openzeppelin/contracts/utils/Address.sol";
+import "@openzeppelin/contracts/token/ERC20/SafeERC20.sol";
+import "@openzeppelin/contracts/math/Math.sol";
 import "../interfaces/BalancerV2.sol";
 import "../interfaces/MasterChef.sol";
 
 contract Strategy is BaseStrategy {
+    using SafeERC20 for IERC20;
+    using Address for address;
+    using SafeMath for uint256;
 
     IBalancerVault public balancerVault;
     IBalancerPool public bpt;
@@ -460,8 +466,7 @@ contract Strategy is BaseStrategy {
         for (uint i = 0; i < rewardTokens.length; i++) {
             rewardTokens[i].approve(address(balancerVault), 0);
         }
-        IERC20[] memory noRewardTokens;
-        rewardTokens = noRewardTokens;
+        delete rewardTokens;
         delete swapSteps;
     }
 
diff --git a/tests/conftest.py b/tests/conftest.py
index c556745..0df92c8 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -130,6 +130,9 @@ def usdc():
 def beets_whale(accounts):
     yield accounts.at("0xa2503804ec837D1E4699932D58a3bdB767DeA505", force=True)
 
+@pytest.fixture
+def tusd_whale(accounts):
+    yield accounts.at("0x789B5DBd47d7Ca3799f8E9FdcE01bC5E356fcDF1", force=True)
 
 @pytest.fixture
 def wftm_amount(user, wftm, accounts):
@@ -203,7 +206,7 @@ def strategyFactory(strategist, keeper, vault, StrategyFactory, gov, balancer_va
                     management,
                     masterChef,
                     swapStepsBeets):
-    factory = strategist.deploy(StrategyFactory, vault, balancer_vault, pool, masterChef, 5, 5, 100_000, 2 * 60 * 60,
+    factory = strategist.deploy(StrategyFactory, vault, balancer_vault, pool, masterChef, 10, 10, 100_000, 2 * 60 * 60,
                                 58)
     yield factory
 
diff --git a/tests/test_operation.py b/tests/test_operation.py
index 3d7ac72..f8fd8a5 100644
--- a/tests/test_operation.py
+++ b/tests/test_operation.py
@@ -70,7 +70,7 @@ def test_manual_exit(
 
 
 def test_profitable_harvest(
-        chain, token, vault, gov, strategy, user, strategist, amount, RELATIVE_APPROX, beets, beets_whale, management
+        chain, token, vault, gov, strategy, user, strategist, amount, RELATIVE_APPROX, beets, beets_whale, management, tusd, tusd_whale
 ):
     # Deposit to the vault
     token.approve(vault.address, amount, {"from": user})
@@ -87,7 +87,7 @@ def test_profitable_harvest(
     chain.mine(1)
     chain.sleep(1)
     # Harvest 2: Realize profit
-    util.airdrop_rewards(strategy, beets, beets_whale)
+    util.airdrop_rewards(strategy, beets, beets_whale, tusd, tusd_whale)
 
     tx = strategy.harvest({"from": strategist})
     print(tx.events["StrategyReported"])
@@ -99,8 +99,50 @@ def test_profitable_harvest(
     assert strategy.estimatedTotalAssets() + profit > amount
     assert vault.pricePerShare() > before_pps
 
+def test_delist_rewards (
+        chain, token, vault, gov, strategy, user, strategist, amount, RELATIVE_APPROX, beets, beets_whale, management, tusd, tusd_whale, swapStepsBeets, swapStepsTusd
+):
+    # Deposit to the vault
+    token.approve(vault.address, amount, {"from": user})
+    vault.deposit(amount, {"from": user})
+    assert token.balanceOf(vault.address) == amount
+
+    # Harvest 1: Send funds through the strategy
+    chain.sleep(1)
+    strategy.harvest({"from": strategist})
+    assert pytest.approx(strategy.estimatedTotalAssets(), rel=RELATIVE_APPROX) == amount
+
+    before_pps = vault.pricePerShare()
+
+    chain.mine(1)
+    chain.sleep(1)
+
+    strategy.delistAllRewards({'from': gov})
+    strategy.whitelistRewards(beets, swapStepsBeets, {'from': gov})
+    strategy.whitelistRewards(tusd, swapStepsTusd, {'from': gov})
+
+    chain.mine(1)
+    chain.sleep(1)
+    # Harvest 2: Realize profit
+    util.airdrop_rewards(strategy, beets, beets_whale, tusd, tusd_whale)
+    assert tusd.balanceOf(strategy) > 0
+    assert beets.balanceOf(strategy) > 0
+    # This should sell rewards
+    tx = strategy.harvest({"from": strategist})
+    print(tx.events["StrategyReported"])
+    chain.sleep(3600 * 6)  # 6 hrs needed for profits to unlock
+    chain.mine(1)
+
+    profit = token.balanceOf(vault.address)
+
+    assert strategy.estimatedTotalAssets() + profit > amount
+    assert vault.pricePerShare() > before_pps
+    assert tusd.balanceOf(strategy) == 0
+    assert beets.balanceOf(strategy) == 0
+
+
 
-def test_deposit_all(chain, token, vault, strategy, user, strategist, amount, RELATIVE_APPROX, beets, beets_whale,
+def test_deposit_all(chain, token, vault, strategy, user, strategist, amount, RELATIVE_APPROX, beets, beets_whale, tusd, tusd_whale,
                      gov):
     # Deposit to the vault
     token.approve(vault.address, amount, {"from": user})
@@ -122,7 +164,7 @@ def test_deposit_all(chain, token, vault, strategy, user, strategist, amount, RE
         chain.mine(1)
 
     before_pps = vault.pricePerShare()
-    util.airdrop_rewards(strategy, beets, beets_whale)
+    util.airdrop_rewards(strategy, beets, beets_whale, tusd, tusd_whale)
     util.stateOfStrat("after airdrop", strategy, beets)
 
     # Harvest 2: Realize profit
@@ -141,7 +183,7 @@ def test_deposit_all(chain, token, vault, strategy, user, strategist, amount, RE
     vault.updateStrategyDebtRatio(strategy.address, 5_000, {"from": gov})
     chain.sleep(1)
 
-    # strategy.setDoHealthCheck(False, {'from':gov})
+    strategy.setDoHealthCheck(False, {'from':gov})
 
     strategy.harvest({"from": strategist})
     util.stateOfStrat("after harvest 5000", strategy, token)
@@ -152,7 +194,7 @@ def test_deposit_all(chain, token, vault, strategy, user, strategist, amount, RE
 
 
 def test_change_debt(
-        chain, gov, token, vault, strategy, user, strategist, amount, RELATIVE_APPROX, beets, beets_whale
+        chain, gov, token, vault, strategy, user, strategist, amount, RELATIVE_APPROX, beets, beets_whale, tusd, tusd_whale
 ):
     # Deposit to the vault and harvest
     token.approve(vault.address, amount, {"from": user})
@@ -171,7 +213,7 @@ def test_change_debt(
     chain.sleep(1)
 
     util.stateOfStrat("before airdrop", strategy, token)
-    util.airdrop_rewards(strategy, beets, beets_whale)
+    util.airdrop_rewards(strategy, beets, beets_whale, tusd, tusd_whale)
     util.stateOfStrat("after airdrop", strategy, token)
 
     vault.updateStrategyDebtRatio(strategy.address, 5_000, {"from": gov})
@@ -185,6 +227,7 @@ def test_change_debt(
 
     vault.updateStrategyDebtRatio(strategy.address, 0, {"from": gov})
     chain.sleep(1)
+    strategy.setDoHealthCheck(False, {'from':gov}) #this harvest will have losses because of slippage
     strategy.harvest({"from": strategist})
     util.stateOfStrat("after harvest", strategy, token)
 
@@ -238,8 +281,8 @@ def test_triggers(
 
 
 # simulate a bad deposit, aka pool has too much of the want you're trying to deposit already
-def test_unbalance_deposit(chain, token, vault, strategy, user, strategist, amount, RELATIVE_APPROX, token2_whale,
-                           token2, token_whale, gov, pool, balancer_vault):
+def test_unbalance_deposit(chain, token, vault, strategy, user, strategist, amount, RELATIVE_APPROX, tusd_whale,
+                           tusd, token_whale, gov, pool, balancer_vault):
     token.approve(vault.address, 2 ** 256 - 1, {"from": user})
     vault.deposit(amount, {"from": user})
     assert token.balanceOf(vault.address) == amount
@@ -251,19 +294,19 @@ def test_unbalance_deposit(chain, token, vault, strategy, user, strategist, amou
 
     print(f'pool rate: {pool.getRate()}')
     tokens = balancer_vault.getPoolTokens(pool.getPoolId())[0]
-    token2Index = 0
-    if (tokens[0] == token2):
-        token2Index = 0
-    elif tokens[1] == token2:
-        token2Index = 1
+    tusdIndex = 0
+    if (tokens[0] == tusd):
+        tusdIndex = 0
+    elif tokens[1] == tusd:
+        tusdIndex = 1
 
-    pooled2 = balancer_vault.getPoolTokens(pool.getPoolId())[1][token2Index]
+    pooled2 = balancer_vault.getPoolTokens(pool.getPoolId())[1][tusdIndex]
     print(balancer_vault.getPoolTokens(pool.getPoolId()))
     print(f'pooled: {pooled2}')
     token.approve(balancer_vault, 2 ** 256 - 1, {'from': token_whale})
 
     # simulate bad pool state by whale to swap out 98% of one side of the pool so pool only has excess want
-    singleSwap = (pool.getPoolId(), 1, token, token2, pooled2 * 0.98, b'0x0')
+    singleSwap = (pool.getPoolId(), 1, token, tusd, pooled2 * 0.98, b'0x0')
     balancer_vault.swap(singleSwap, (token_whale, False, token_whale, False), token.balanceOf(token_whale),
                         2 ** 256 - 1, {'from': token_whale})
     print(balancer_vault.getPoolTokens(pool.getPoolId()))
@@ -278,7 +321,7 @@ def test_unbalance_deposit(chain, token, vault, strategy, user, strategist, amou
 
 
 def test_unbalanced_pool_withdraw(chain, token, vault, strategy, user, strategist, amount, RELATIVE_APPROX,
-                                  token2_whale, token2,
+                                  tusd_whale, tusd,
                                   gov, pool, balancer_vault):
     # Deposit to the vault
     token.approve(vault.address, amount, {"from": user})
@@ -305,22 +348,22 @@ def test_unbalanced_pool_withdraw(chain, token, vault, strategy, user, strategis
 
     print(f'pool rate: {pool.getRate()}')
     tokens = balancer_vault.getPoolTokens(pool.getPoolId())[0]
-    token2Index = 0
-    if (tokens[0] == token2):
-        token2Index = 0
-    elif tokens[1] == token2:
-        token2Index = 1
+    tusdIndex = 0
+    if (tokens[0] == tusd):
+        tusdIndex = 0
+    elif tokens[1] == tusd:
+        tusdIndex = 1
 
     util.stateOfStrat("after deposit all    ", strategy, token)
     pooled = balancer_vault.getPoolTokens(pool.getPoolId())[1][0]
     print(balancer_vault.getPoolTokens(pool.getPoolId()))
     print(f'pooled: {pooled}')
-    token2.approve(balancer_vault, 2 ** 256 - 1, {'from': token2_whale})
+    tusd.approve(balancer_vault, 2 ** 256 - 1, {'from': tusd_whale})
 
     # simulate bad pool state by whale to swap out 98% of one side of the pool so pool only has 2% of the original want
-    singleSwap = (pool.getPoolId(), 1, token2, token, pooled * 0.98, b'0x0')
-    balancer_vault.swap(singleSwap, (token2_whale, False, token2_whale, False), token2.balanceOf(token2_whale),
-                        2 ** 256 - 1, {'from': token2_whale})
+    singleSwap = (pool.getPoolId(), 1, tusd, token, pooled * 0.98, b'0x0')
+    balancer_vault.swap(singleSwap, (tusd_whale, False, tusd_whale, False), tusd.balanceOf(tusd_whale),
+                        2 ** 256 - 1, {'from': tusd_whale})
     print(balancer_vault.getPoolTokens(pool.getPoolId()))
     print(f'pool rate: {pool.getRate()}')
 
diff --git a/tests/util.py b/tests/util.py
index cac9efa..2e51097 100644
--- a/tests/util.py
+++ b/tests/util.py
@@ -1,19 +1,17 @@
 from brownie import Contract
 
 
-def airdrop_rewards(strategy, beets, beets_whale):
+def airdrop_beets_rewards(strategy, beets, beets_whale):
     beets.approve(strategy, 2 ** 256 - 1, {'from': beets_whale})
     beets.transfer(strategy, 500 * 1e18, {'from': beets_whale})
 
-def airdrop_tusd_rewards(strategy, token2, token2_whale):
-    token2.approve(strategy, 2 ** 256 - 1, {'from': token2_whale})
-    token2.transfer(strategy, 3_000 * 1e18, {'from': token2_whale})
+def airdrop_tusd_rewards(strategy, tusd, tusd_whale):
+    tusd.approve(strategy, 2 ** 256 - 1, {'from': tusd_whale})
+    tusd.transfer(strategy, 500 * 1e18, {'from': tusd_whale})
 
-def airdrop_all_rewards(strategy, beets, beets_whale, token2, token2_whale):
-    token2.approve(strategy, 2 ** 256 - 1, {'from': token2_whale})
-    token2.transfer(strategy, 3_000 * 1e18, {'from': token2_whale})
-    token2.approve(strategy, 2 ** 256 - 1, {'from': token2_whale})
-    token2.transfer(strategy, 3_000 * 1e18, {'from': token2_whale})
+def airdrop_rewards(strategy, beets, beets_whale, tusd, tusd_whale):
+    airdrop_beets_rewards(strategy, beets, beets_whale)
+    airdrop_tusd_rewards(strategy, tusd, tusd_whale)
 
 
 def stateOfStrat(msg, strategy, token):

From 8c0de62c7e47c29405efb29e902a2d82121fae5d Mon Sep 17 00:00:00 2001
From: Micky Mousse <newmickymousse@pm.me>
Date: Tue, 29 Mar 2022 22:58:25 +0200
Subject: [PATCH 4/5] fix: security review's comments

---
 contracts/Strategy.sol  |  1 +
 tests/conftest.py       |  4 +++-
 tests/test_operation.py | 11 +++++++----
 3 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/contracts/Strategy.sol b/contracts/Strategy.sol
index 5afb0a5..685ede7 100644
--- a/contracts/Strategy.sol
+++ b/contracts/Strategy.sol
@@ -456,6 +456,7 @@ contract Strategy is BaseStrategy {
 
     // for partnership rewards like TUSD or airdrops
     function whitelistRewards(address _rewardToken, SwapSteps memory _steps) public isVaultManager {
+        require(address(_rewardToken) != address(want));
         IERC20 token = IERC20(_rewardToken);
         token.approve(address(balancerVault), max);
         rewardTokens.push(token);
diff --git a/tests/conftest.py b/tests/conftest.py
index 0df92c8..1131e63 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -75,7 +75,7 @@ def token_whale(accounts):
     # 0x93C08a3168fC469F3fC165cd3A471D19a37ca19e USDC
     # 0x8D9AED9882b4953a0c9fa920168fa1FDfA0eBE75 DAI
     # 0x789B5DBd47d7Ca3799f8E9FdcE01bC5E356fcDF1 TUSD
-    return accounts.at("0x93C08a3168fC469F3fC165cd3A471D19a37ca19e", force=True)
+    return accounts.at("0xc5ed2333f8a2c351fca35e5ebadb2a82f5d254c3", force=True)
 
 
 @pytest.fixture
@@ -221,6 +221,8 @@ def strategy(strategist, keeper, vault, strategyFactory, gov, balancer_vault, po
     vault.addStrategy(strategy, 10_000, 0, 2 ** 256 - 1, 1_000, {"from": gov})
     vault.setManagementFee(0, {"from": gov})
     chain.sleep(1)
+    strategy.setParams(25, 25, strategy.maxSingleDeposit(), strategy.minDepositPeriod(),
+            strategy.keep(), strategy.keepBips(), {"from": gov})
     yield strategy
 
 
diff --git a/tests/test_operation.py b/tests/test_operation.py
index f8fd8a5..0e7d8be 100644
--- a/tests/test_operation.py
+++ b/tests/test_operation.py
@@ -5,7 +5,7 @@
 
 
 def test_operation(
-        chain, accounts, token, vault, strategy, user, strategist, amount, RELATIVE_APPROX
+        chain, accounts, token, vault, strategy, user, strategist, amount, gov, RELATIVE_APPROX
 ):
     # Deposit to the vault
     print("Strategy Name:", strategy.name())
@@ -22,9 +22,11 @@ def test_operation(
     # tend()
     strategy.tend({"from": strategist})
 
+    chain.sleep(1)
     # withdrawal
-    vault.withdraw(vault.balanceOf(user), user, 10, {"from": user})
-    assert (pytest.approx(token.balanceOf(user), rel=RELATIVE_APPROX) == user_balance_before)
+    vault.withdraw(vault.balanceOf(user), user, 20, {"from": user})
+    assert (token.balanceOf(user) >= user_balance_before * 0.98) #2% loss
+    assert (token.balanceOf(user) <= user_balance_before)
 
 
 def test_emergency_exit(
@@ -204,7 +206,8 @@ def test_change_debt(
     strategy.harvest({"from": strategist})
     half = int(amount / 2)
 
-    assert pytest.approx(strategy.estimatedTotalAssets(), rel=RELATIVE_APPROX) == half
+    assert pytest.approx(strategy.estimatedTotalAssets(),
+                         rel=RELATIVE_APPROX) == half or strategy.estimatedTotalAssets() >= half
 
     vault.updateStrategyDebtRatio(strategy.address, 10_000, {"from": gov})
     chain.sleep(1)

From 1fa1e0321931e2d96e9b6bbfbc5c06e791f756af Mon Sep 17 00:00:00 2001
From: Micky Mousse <newmickymousse@pm.me>
Date: Fri, 22 Apr 2022 12:18:27 +0200
Subject: [PATCH 5/5] fix: security review's comments

---
 tests/test_shutdown.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tests/test_shutdown.py b/tests/test_shutdown.py
index 34d12ac..e85b1db 100644
--- a/tests/test_shutdown.py
+++ b/tests/test_shutdown.py
@@ -23,9 +23,10 @@ def test_vault_shutdown_can_withdraw(
     vault.setEmergencyShutdown(True)
 
     ## Withdraw (does it work, do you get what you expect)
-    vault.withdraw(vault.balanceOf(user), user, 10, {"from": user})
+    vault.withdraw(vault.balanceOf(user), user, 20, {"from": user})
 
-    assert pytest.approx(token.balanceOf(user), rel=RELATIVE_APPROX) == amount
+    assert token.balanceOf(user) >= amount * 0.98
+    assert token.balanceOf(user) <= amount
 
 
 def test_basic_shutdown(