fix: update to strategy changes

yearn · Jan 29, 2024 · 4d27d31 · 4d27d31
1 parent d39e2ae
commit 4d27d31
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 23 deletions.
diff --git a/contracts/test/mocks/ERC4626/LossyStrategy.sol b/contracts/test/mocks/ERC4626/LossyStrategy.sol
@@ -41,13 +41,13 @@ contract ERC4626LossyStrategy is MockTokenizedStrategy {
         yieldSource = address(new YieldSource(_asset));
         ERC20(_asset).safeApprove(yieldSource, type(uint256).max);
         // So we can record losses when it happens.
-        strategyStorage().management = address(this);
+        _strategyStorage().management = address(this);
         vault = _vault;
     }
 
     // used to generate losses, accepts single arg to send losses to
     function setLoss(address _target, uint256 _loss) external {
-        strategyStorage().asset.safeTransferFrom(yieldSource, _target, _loss);
+        _strategyStorage().asset.safeTransferFrom(yieldSource, _target, _loss);
         // Record the loss
         MockTokenizedStrategy(address(this)).report();
     }
@@ -71,7 +71,7 @@ contract ERC4626LossyStrategy is MockTokenizedStrategy {
 
         if (withdrawingLoss < 0) {
             // Over withdraw to the vault
-            strategyStorage().asset.safeTransfer(
+            _strategyStorage().asset.safeTransfer(
                 vault,
                 uint256(-withdrawingLoss)
             );
@@ -80,16 +80,16 @@ contract ERC4626LossyStrategy is MockTokenizedStrategy {
 
     function harvestAndReport() external override returns (uint256) {
         return
-            strategyStorage().asset.balanceOf(address(this)) +
-            strategyStorage().asset.balanceOf(yieldSource);
+            _strategyStorage().asset.balanceOf(address(this)) +
+            _strategyStorage().asset.balanceOf(yieldSource);
     }
 
     function availableWithdrawLimit(
         address
     ) public view override returns (uint256) {
         return
-            strategyStorage().asset.balanceOf(address(this)) +
-            strategyStorage().asset.balanceOf(yieldSource) -
+            _strategyStorage().asset.balanceOf(address(this)) +
+            _strategyStorage().asset.balanceOf(yieldSource) -
             lockedFunds;
     }
 }
diff --git a/contracts/test/mocks/ERC4626/MockTokenizedStrategy.sol b/contracts/test/mocks/ERC4626/MockTokenizedStrategy.sol
@@ -7,27 +7,14 @@ contract MockTokenizedStrategy is TokenizedStrategy {
     uint256 public minDebt;
     uint256 public maxDebt = type(uint256).max;
 
-    // Private variables and functions used in this mock.
-    bytes32 public constant BASE_STRATEGY_STORAGE =
-        bytes32(uint256(keccak256("yearn.base.strategy.storage")) - 1);
-
-    function strategyStorage() internal pure returns (StrategyData storage S) {
-        // Since STORAGE_SLOT is a constant, we have to put a variable
-        // on the stack to access it from an inline assembly block.
-        bytes32 slot = BASE_STRATEGY_STORAGE;
-        assembly {
-            S.slot := slot
-        }
-    }
-
     constructor(
         address _asset,
         string memory _name,
         address _management,
         address _keeper
     ) {
         // Cache storage pointer
-        StrategyData storage S = strategyStorage();
+        StrategyData storage S = _strategyStorage();
 
         // Set the strategy's underlying asset
         S.asset = ERC20(_asset);
@@ -58,7 +45,7 @@ contract MockTokenizedStrategy is TokenizedStrategy {
     function availableDepositLimit(
         address
     ) public view virtual returns (uint256) {
-        uint256 _totalAssets = strategyStorage().totalAssets;
+        uint256 _totalAssets = _strategyStorage().totalAssets;
         uint256 _maxDebt = maxDebt;
         return _maxDebt > _totalAssets ? _maxDebt - _totalAssets : 0;
     }
@@ -74,6 +61,6 @@ contract MockTokenizedStrategy is TokenizedStrategy {
     function freeFunds(uint256 _amount) external virtual {}
 
     function harvestAndReport() external virtual returns (uint256) {
-        return strategyStorage().asset.balanceOf(address(this));
+        return _strategyStorage().asset.balanceOf(address(this));
     }
 }