decaf-addict · decaf-addict · Sep 30, 2021 · Oct 1, 2021 · fp-crypto · Oct 19, 2021
diff --git a/contracts/Strategy.sol b/contracts/Strategy.sol
@@ -53,6 +53,7 @@ contract Strategy is BaseStrategy {
     uint256 public lastDepositTime;
     uint256 internal constant basisOne = 10000;
     bool internal isOriginal = true;
+    uint internal etaCached;
 
     constructor(
         address _vault,
@@ -160,18 +161,28 @@ contract Strategy is BaseStrategy {
     }
 
     function estimatedTotalAssets() public view override returns (uint256) {
-        return balanceOfWant().add(balanceOfPooled());
+        return etaCached;
+    }
+
+    // There is no way to calculate the total asset without doing a tx call.
+    function estimateTotalAssets() public returns (uint256 _wants) {
+        etaCached = balanceOfWant().add(balanceOfPooled());
+        return etaCached;
     }
 
     function prepareReturn(uint256 _debtOutstanding) internal override returns (uint256 _profit, uint256 _loss, uint256 _debtPayment){
+        uint256 total = estimateTotalAssets();
+        uint256 debt = vault.strategies(address(this)).totalDebt;
+        uint toCollect = total > debt ? total.sub(debt) : 0;
+
         if (_debtOutstanding > 0) {
             (_debtPayment, _loss) = liquidatePosition(_debtOutstanding);
         }
 
         uint256 beforeWant = balanceOfWant();
 
         // 2 forms of profit. Incentivized rewards (BAL+other) and pool fees (want)
-        collectTradingFees();
+        collectTradingFees(toCollect);
         sellRewards();
 
         uint256 afterWant = balanceOfWant();
@@ -213,7 +224,7 @@ contract Strategy is BaseStrategy {
     }
 
     function liquidatePosition(uint256 _amountNeeded) internal override returns (uint256 _liquidatedAmount, uint256 _loss){
-        if (estimatedTotalAssets() < _amountNeeded) {
+        if (estimateTotalAssets() < _amountNeeded) {
             _liquidatedAmount = liquidateAllPositions();
             return (_liquidatedAmount, _amountNeeded.sub(_liquidatedAmount));
         }
@@ -234,7 +245,7 @@ contract Strategy is BaseStrategy {
     }
 
     function liquidateAllPositions() internal override returns (uint256 liquidated) {
-        uint eta = estimatedTotalAssets();
+        uint eta = estimateTotalAssets();
         uint256 bpts = balanceOfBpt();
         if (bpts > 0) {
             // exit entire position for single token. Could revert due to single exit limit enforced by balancer
@@ -315,12 +326,9 @@ contract Strategy is BaseStrategy {
         }
     }
 
-    function collectTradingFees() internal {
-        uint256 total = estimatedTotalAssets();
-        uint256 debt = vault.strategies(address(this)).totalDebt;
-        if (total > debt) {
-            uint256 profit = total.sub(debt);
-            _sellBptForExactToken(profit);
+    function collectTradingFees(uint _profit) internal {
+        if (_profit > 0) {
+            _sellBptForExactToken(_profit);
         }
     }
 
@@ -336,19 +344,37 @@ contract Strategy is BaseStrategy {
         return rewardTokens[index].balanceOf(address(this));
     }
 
-    function balanceOfPooled() public view returns (uint256 _amount){
+    function balanceOfPooled() public returns (uint256 _amount){
         uint256 totalWantPooled;
-        (IERC20[] memory tokens,uint256[] memory totalBalances,uint256 lastChangeBlock) = balancerVault.getPoolTokens(balancerPoolId);
-        for (uint8 i = 0; i < numTokens; i++) {
-            uint256 tokenPooled = totalBalances[i].mul(balanceOfBpt()).div(bpt.totalSupply());
-            if (tokenPooled > 0) {
-                IERC20 token = tokens[i];
-                if (token != want) {
-                    IBalancerPool.SwapRequest memory request = _getSwapRequest(token, tokenPooled, lastChangeBlock);
-                    // now denomated in want
-                    tokenPooled = bpt.onSwap(request, totalBalances, i, tokenIndex);
+        uint bpts = balanceOfBpt();
+        if (bpts > 0) {
+            (IERC20[] memory tokens,uint256[] memory totalBalances,uint256 lastChangeBlock) = balancerVault.getPoolTokens(balancerPoolId);
+            for (uint8 i = 0; i < numTokens; i++) {
+                uint256 tokenPooled = totalBalances[i].mul(bpts).div(bpt.totalSupply());
+
+                if (tokenPooled > 0) {
+                    if (tokens[i] != want) {
+                        // single step bc doing one swap within own pool i.e. wsteth -> weth
+                        IBalancerVault.BatchSwapStep[] memory steps = new IBalancerVault.BatchSwapStep[](1);
+                        steps[0] = IBalancerVault.BatchSwapStep(balancerPoolId,
+                            tokenIndex == 0 ? 1 : 0,
+                            tokenIndex,
+                            tokenPooled,
+                            abi.encode(0)
+                        );
+
+                        // 2 assets of the pool, ordered by trade direction i.e. wsteth -> weth
+                        IAsset[] memory path = new IAsset[](2);
+                        path[0] = IAsset(address(tokenIndex == 0 ? tokens[1] : tokens[0]));
+                        path[1] = IAsset(address(want));
+
+                        tokenPooled = uint(- balancerVault.queryBatchSwap(IBalancerVault.SwapKind.GIVEN_IN,
+                            steps,
+                            path,
+                            IBalancerVault.FundManagement(address(this), false, address(this), false))[1]);
+                    }
+                    totalWantPooled += tokenPooled;
                 }
-                totalWantPooled += tokenPooled;
             }
         }
         return totalWantPooled;

diff --git a/interfaces/BalancerV2.sol b/interfaces/BalancerV2.sol
@@ -29,7 +29,7 @@ interface IBalancerPool is IERC20 {
         uint256[] memory balances,
         uint256 indexIn,
         uint256 indexOut
-    ) external view returns (uint256 amount);
+    ) external returns (uint256 amount);
 }
 
 interface IBalancerVault {
@@ -198,6 +198,14 @@ interface IBalancerVault {
         int256[] memory limits,
         uint256 deadline
     ) external payable returns (int256[] memory);
+
+    // CAVEAT!! Do not call this after a batchSwap in the same txn
+    function queryBatchSwap(
+        SwapKind kind,
+        BatchSwapStep[] memory swaps,
+        IAsset[] memory assets,
+        FundManagement memory funds
+    ) external returns (int256[] memory);
 }
 
 interface IAsset {

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -50,7 +50,7 @@ def token():
     # 0xdAC17F958D2ee523a2206206994597C13D831ec7 USDT
     # 0x7f39C581F595B53c5cb19bD0b3f8dA6c935E2Ca0 wSTETH
     # 0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2 WETH
-    token_address = "0x6B175474E89094C44Da98b954EedeAC495271d0F"
+    token_address = "0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2"
     yield Contract(token_address)
 
 
@@ -61,7 +61,7 @@ def token2():
     # 0xdAC17F958D2ee523a2206206994597C13D831ec7 USDT
     # 0x7f39C581F595B53c5cb19bD0b3f8dA6c935E2Ca0 wSTETH
     # 0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2 WETH
-    token_address = "0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48"
+    token_address = "0x7f39C581F595B53c5cb19bD0b3f8dA6c935E2Ca0"
     yield Contract(token_address)
 
 
@@ -72,25 +72,27 @@ def token_whale(accounts):
     # 0xA929022c9107643515F5c777cE9a910F0D1e490C USDT
     # 0xba12222222228d8ba445958a75a0704d566bf2c8 wSTETH
     # 0x2F0b23f53734252Bda2277357e97e1517d6B042A WETH
-    return accounts.at("0x5d3a536E4D6DbD6114cc1Ead35777bAB948E3643", force=True)
+    return accounts.at("0x2F0b23f53734252Bda2277357e97e1517d6B042A", force=True)
 
 
 @pytest.fixture
 def amount(accounts, token, user, token_whale):
-    amount = 10_000_000 * 10 ** token.decimals()
+    amount = 300 * 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
 
 
 @pytest.fixture
 def amount2(accounts, token2, user):
-    amount = 1_000_000 * 10 ** token2.decimals()
+    amount = 300 * 10 ** token2.decimals()
     # In order to get some funds for the token you are about to use,
     # 0x5d3a536E4D6DbD6114cc1Ead35777bAB948E3643 DAI
     # 0x0A59649758aa4d66E25f08Dd01271e891fe52199 USDC
     # 0xA929022c9107643515F5c777cE9a910F0D1e490C USDT
-    reserve = accounts.at("0x0A59649758aa4d66E25f08Dd01271e891fe52199", force=True)
+    # 0xba12222222228d8ba445958a75a0704d566bf2c8 wSTETH
+    # 0x2F0b23f53734252Bda2277357e97e1517d6B042A WETH
+    reserve = accounts.at("0xba12222222228d8ba445958a75a0704d566bf2c8", force=True)
     token2.transfer(user, amount, {"from": reserve})
     yield amount
 
@@ -101,6 +103,12 @@ def weth():
     yield Contract(token_address)
 
 
+@pytest.fixture
+def wsteth():
+    token_address = "0x7f39C581F595B53c5cb19bD0b3f8dA6c935E2Ca0"
+    yield Contract(token_address)
+
+
 @pytest.fixture
 def bal():
     token_address = "0xba100000625a3754423978a60c9317c58a424e3D"
@@ -171,52 +179,58 @@ def balancer_vault():
 def pool():
     # 0x06Df3b2bbB68adc8B0e302443692037ED9f91b42 stable pool
     # 0x32296969Ef14EB0c6d29669C550D4a0449130230 metastable eth pool
-    address = "0x06Df3b2bbB68adc8B0e302443692037ED9f91b42"
+    address = "0x32296969Ef14EB0c6d29669C550D4a0449130230"
     yield Contract(address)
 
 
 @pytest.fixture
 def balWethPoolId():
-    yield 0x5c6ee304399dbdb9c8ef030ab642b10820db8f56000200000000000000000014
+    yield 0x5c6ee304399dbdb9c8ef030ab642b10820db8f56000200000000000000000014  # bal-weth
 
 
 @pytest.fixture
 def wethTokenPoolId():
     id = 0x0b09dea16768f0799065c475be02919503cb2a3500020000000000000000001a  # weth-dai
     yield id
 
+
 @pytest.fixture
 def wethToken2PoolId():
-    id = 0x96646936b91d6b9d7d0c47c496afbf3d6ec7b6f8000200000000000000000019  # weth-usdc
+    # id = 0x96646936b91d6b9d7d0c47c496afbf3d6ec7b6f8000200000000000000000019  # weth-usdc
+    id = 0x32296969ef14eb0c6d29669c550d4a0449130230000200000000000000000080  # weth-wsteth
     yield id
 
+
 @pytest.fixture
 def ldoWethPoolId():
-    id = 0xbf96189eee9357a95c7719f4f5047f76bde804e5000200000000000000000087  # weth-dai
+    id = 0xbf96189eee9357a95c7719f4f5047f76bde804e5000200000000000000000087  # ldo-weth
     yield id
 
 
 @pytest.fixture
-def swapStepsBal(balWethPoolId, wethTokenPoolId, bal, weth, token):
-    yield ([balWethPoolId, wethTokenPoolId], [bal, weth, token])
+def swapStepsBal(balWethPoolId, bal, weth):
+    yield ([balWethPoolId], [bal, weth])
+
 
 @pytest.fixture
-def swapStepsLdo(ldoWethPoolId, wethTokenPoolId, ldo, weth, token):
-    yield ([ldoWethPoolId, wethTokenPoolId], [ldo, weth, token])
+def swapStepsLdo(ldoWethPoolId, ldo, weth):
+    yield ([ldoWethPoolId], [ldo, weth])
+
 
 @pytest.fixture
-def swapStepsBal2(balWethPoolId, wethToken2PoolId, bal, weth, token2):
-    yield ([balWethPoolId, wethToken2PoolId], [bal, weth, token2])
+def swapStepsBal2(balWethPoolId, wethToken2PoolId, bal, weth, wsteth):
+    yield ([balWethPoolId], [bal, weth, wsteth])
+
 
 @pytest.fixture
-def swapStepsLdo2(ldoWethPoolId, wethToken2PoolId, ldo, weth, token2):
-    yield ([ldoWethPoolId, wethToken2PoolId], [ldo, weth, token2])
+def swapStepsLdo2(ldoWethPoolId, wethToken2PoolId, ldo, weth, wsteth):
+    yield ([ldoWethPoolId], [ldo, weth, wsteth])
 
 
 @pytest.fixture
 def strategy(strategist, keeper, vault, Strategy, gov, balancer_vault, pool, bal, ldo, management, swapStepsBal,
              swapStepsLdo):
-    strategy = strategist.deploy(Strategy, vault, balancer_vault, pool, 5, 5, 1_000_000, 2 * 60 * 60)
+    strategy = strategist.deploy(Strategy, vault, balancer_vault, pool, 5, 5, 30, 2 * 60 * 60)
     strategy.setKeeper(keeper)
     strategy.whitelistRewards(bal, swapStepsBal, {'from': gov})
     strategy.whitelistRewards(ldo, swapStepsLdo, {'from': gov})

diff --git a/tests/test_clone.py b/tests/test_clone.py
@@ -4,28 +4,30 @@
 import test_operation
 
 
-def test_clone(accounts, Strategy, strategy, strategist, rewards, keeper, token2, user, vault, vault2, amount2,
-               balancer_vault,
-               pool, chain, gov,
-               RELATIVE_APPROX,
-               bal, bal_whale, ldo, weth, weth_amout,
-               ldo_whale, swapStepsBal2, swapStepsLdo2, management):
-    with brownie.reverts("Strategy already initialized"):
-        strategy.initialize(vault, strategist, rewards, keeper, balancer_vault, pool, 10, 10, 100_000, 2 * 60 * 60)
+# unique pool. No clones
 
-    transaction = strategy.clone(vault2, strategist, rewards, keeper, balancer_vault, pool, 10, 10, 100_000,
-                                 2 * 60 * 60)
-    cloned_strategy = Strategy.at(transaction.return_value)
-
-    with brownie.reverts("Strategy already initialized"):
-        cloned_strategy.initialize(vault, strategist, rewards, keeper, balancer_vault, pool, 10, 10, 100_000,
-                                   2 * 60 * 60, {'from': gov})
-    cloned_strategy.setKeeper(keeper, {'from': gov})
-    cloned_strategy.whitelistRewards(bal, swapStepsBal2, {'from': management})
-    cloned_strategy.whitelistRewards(ldo, swapStepsLdo2, {'from': management})
-    vault2.addStrategy(cloned_strategy, 10_000, 0, 2 ** 256 - 1, 1_000, {"from": gov})
-
-    # test operations with clone strategy
-    test_operation.test_profitable_harvest(
-        chain, token2, vault2, cloned_strategy, user, strategist, amount2, RELATIVE_APPROX, bal, bal_whale, ldo,
-        ldo_whale, management)
+# def test_clone(accounts, Strategy, strategy, strategist, rewards, keeper, token2, user, vault, vault2, amount2,
+#                balancer_vault,
+#                pool, chain, gov,
+#                RELATIVE_APPROX,
+#                bal, bal_whale, ldo, weth, weth_amout,
+#                ldo_whale, swapStepsBal2, swapStepsLdo2, management):
+#     with brownie.reverts("Strategy already initialized"):
+#         strategy.initialize(vault, strategist, rewards, keeper, balancer_vault, pool, 10, 10, 100_000, 2 * 60 * 60)
+#
+#     transaction = strategy.clone(vault2, strategist, rewards, keeper, balancer_vault, pool, 10, 10, 100_000,
+#                                  2 * 60 * 60)
+#     cloned_strategy = Strategy.at(transaction.return_value)
+#
+#     with brownie.reverts("Strategy already initialized"):
+#         cloned_strategy.initialize(vault, strategist, rewards, keeper, balancer_vault, pool, 10, 10, 100_000,
+#                                    2 * 60 * 60, {'from': gov})
+#     cloned_strategy.setKeeper(keeper, {'from': gov})
+#     cloned_strategy.whitelistRewards(bal, swapStepsBal2, {'from': management})
+#     cloned_strategy.whitelistRewards(ldo, swapStepsLdo2, {'from': management})
+#     vault2.addStrategy(cloned_strategy, 10_000, 0, 2 ** 256 - 1, 1_000, {"from": gov})
+#
+#     # test operations with clone strategy
+#     test_operation.test_profitable_harvest(
+#         chain, token2, vault2, cloned_strategy, user, strategist, amount2, RELATIVE_APPROX, bal, bal_whale, ldo,
+#         ldo_whale, management)
diff --git a/tests/test_migration.py b/tests/test_migration.py
@@ -20,9 +20,9 @@ def test_migration(
     vault.deposit(amount, {"from": user})
     chain.sleep(1)
     strategy.harvest()
-    assert pytest.approx(strategy.estimatedTotalAssets(), rel=RELATIVE_APPROX) == amount
+    assert pytest.approx(strategy.estimateTotalAssets({"from": user}).return_value, rel=RELATIVE_APPROX) == amount
 
     # migrate to a new strategy
     new_strategy = strategist.deploy(Strategy, vault, balancer_vault, pool, 5, 5, 100_000, 2 * 60 * 60)
     vault.migrateStrategy(strategy, new_strategy, {"from": gov})
-    assert (pytest.approx(new_strategy.estimatedTotalAssets(), rel=RELATIVE_APPROX) == amount)
+    assert (pytest.approx(new_strategy.estimateTotalAssets({"from": user}).return_value, rel=RELATIVE_APPROX) == amount)