Enable pool fees when rebalancing Aerodrome AMO

contracts/contracts/strategies/aerodrome/AerodromeAMOStrategy.sol

@@ -429,8 +429,9 @@
          * liquidity in the pool yet. The full liquidity removal is thus skipped.
          */
         if (tokenId != 0) {
-            _removeLiquidity(1e18);
+            _removeLiquidityToFacilitateSwap(_amountToSwap, _swapWeth);
         }
+
         // in some cases we will just want to add liquidity and not issue a swap to move the
         // active trading position within the pool
         if (_amountToSwap > 0) {
@@ -472,6 +473,28 @@
         }
     }
 
+    /**
+     * @dev Remove just enough liquidity so that the swap operation can be performed. No need to
+     * add the gaugeUnstakeAndRestake modifier since the underlying _removeLiquidity function
+     * will do that.
+     * Good to consider that remaining liquidity will remain in the pool & staked to the gauge.
+     *
+     * @param _amountToSwap The amount of the token to swap
+     * @param _swapWeth Swap using WETH when true, use OETHb when false
+     */
+    function _removeLiquidityToFacilitateSwap(
+        uint256 _amountToSwap,
+        bool _swapWeth
+    ) internal {
+        // swapping OETHb to WETH doesn't require liquidity removal or when there 
+        // is no amount to be swapped
+        if (!_swapWeth || _amountToSwap == 0) {
+            return;
+        }
+
+        _assureWETHBalance(_amountToSwap);
+    }
+
     /**
      * @dev Decrease partial or all liquidity from the pool.
      * @param _liquidityToDecrease The amount of liquidity to remove expressed in 18 decimal point
@@ -545,6 +568,8 @@
         uint256 _balance = _tokenToSwap.balanceOf(address(this));
 
         if (_balance < _amountToSwap) {
+            // This should never trigger since _removeLiquidityToFacilitateSwap will already
+            // throw an error if there is not enough WETH
             if (_swapWeth) {
                 revert NotEnoughWethForSwap(_balance, _amountToSwap);
             }
@@ -590,7 +615,9 @@
     function _addLiquidity() internal gaugeUnstakeAndRestake {
         uint256 _wethBalance = IERC20(WETH).balanceOf(address(this));
         uint256 _oethbBalance = IERC20(OETHb).balanceOf(address(this));
-        require(_wethBalance > 0, "Must add some WETH");
+        if (_wethBalance == 0) {
+            return;
+        }
 
         uint160 _currentPrice = getPoolX96Price();
         /**
@@ -803,20 +830,12 @@
     }
 
     /**
-     * @notice Withdraw an `amount` of assets from the platform and
-     *         send to the `_recipient`.
-     * @param _recipient  Address to which the asset should be sent
-     * @param _asset      WETH address
-     * @param _amount     Amount of WETH to withdraw
+     * @dev This function removes the appropriate amount of liquidity to assure that the required
+     * amount of WETH is available on the contract
+     *
+     * @param _amount  WETH balance required on the contract
      */
-    function withdraw(
-        address _recipient,
-        address _asset,
-        uint256 _amount
-    ) external override onlyVault nonReentrant {
-        require(_asset == WETH, "Unsupported asset");
-        require(_recipient == vaultAddress, "Only withdraw to vault allowed");
-
+    function _assureWETHBalance(uint256 _amount) internal {
         uint256 _wethBalance = IERC20(WETH).balanceOf(address(this));
         if (_wethBalance < _amount) {
             require(tokenId != 0, "No liquidity available");
@@ -836,6 +855,24 @@
             );
             _removeLiquidity(shareOfWethToRemove);
         }
+    }
+
+    /**
+     * @notice Withdraw an `amount` of assets from the platform and
+     *         send to the `_recipient`.
+     * @param _recipient  Address to which the asset should be sent
+     * @param _asset      WETH address
+     * @param _amount     Amount of WETH to withdraw
+     */
+    function withdraw(
+        address _recipient,
+        address _asset,
+        uint256 _amount
+    ) external override onlyVault nonReentrant {
+        require(_asset == WETH, "Unsupported asset");
+        require(_recipient == vaultAddress, "Only withdraw to vault allowed");
+
+        _assureWETHBalance(_amount);
 
         // burn remaining OETHb
         _burnOethbOnTheContract();

contracts/deploy/base/018_upgrade_amo.js

@@ -0,0 +1,28 @@
+const { deployOnBaseWithGuardian } = require("../../utils/deploy-l2");
+const {
+  deployBaseAerodromeAMOStrategyImplementation,
+} = require("../deployActions");
+
+module.exports = deployOnBaseWithGuardian(
+  {
+    deployName: "018_upgrade_amo",
+  },
+  async ({ ethers }) => {
+    const cAMOStrategyProxy = await ethers.getContract(
+      "AerodromeAMOStrategyProxy"
+    );
+    const cAMOStrategyImpl =
+      await deployBaseAerodromeAMOStrategyImplementation();
+
+    return {
+      actions: [
+        {
+          // 1. Upgrade AMO
+          contract: cAMOStrategyProxy,
+          signature: "upgradeTo(address)",
+          args: [cAMOStrategyImpl.address],
+        },
+      ],
+    };
+  }
+);

contracts/test/_fixture-base.js

@@ -148,14 +148,15 @@ const defaultBaseFixture = deployments.createFixture(async () => {
   );
   const oethVaultSigner = await impersonateAccount(oethbVault.address);
 
-  let strategist;
+  let strategist, harvesterSigner;
   if (isFork) {
     // Impersonate strategist on Fork
     strategist = await impersonateAndFund(strategistAddr);
     strategist.address = strategistAddr;
 
     await impersonateAndFund(governor.address);
     await impersonateAndFund(timelock.address);
+    harvesterSigner = await impersonateAndFund(harvester.address);
 
     // configure Vault to not automatically deposit to strategy
     await oethbVault.connect(governor).setVaultBuffer(oethUnits("1"));
@@ -208,6 +209,7 @@ const defaultBaseFixture = deployments.createFixture(async () => {
     wOETHb,
     zapper,
     harvester,
+    harvesterSigner,
     dripper,
 
     // Bridged WOETH

contracts/test/strategies/aerodrome-amo.base.fork-test.js

@@ -270,6 +270,8 @@ describe("ForkTest: Aerodrome AMO Strategy (Base)", async function () {
     rafael,
     aeroSwapRouter,
     aeroNftManager,
+    harvester,
+    harvesterSigner,
     quoter;
 
   beforeEach(async () => {
@@ -286,6 +288,8 @@ describe("ForkTest: Aerodrome AMO Strategy (Base)", async function () {
     aeroNftManager = fixture.aeroNftManager;
     oethbVaultSigner = await impersonateAndFund(oethbVault.address);
     gauge = fixture.aeroClGauge;
+    harvester = fixture.harvester;
+    harvesterSigner = fixture.harvesterSigner;
     quoter = fixture.quoter;
 
     await setup();
@@ -369,7 +373,7 @@ describe("ForkTest: Aerodrome AMO Strategy (Base)", async function () {
 
       // correct harvester set
       expect(await aerodromeAmoStrategy.harvesterAddress()).to.equal(
-        await strategist.getAddress()
+        addresses.base.HarvesterProxy
       );
 
       await assetLpStakedInGauge();
@@ -471,18 +475,16 @@ describe("ForkTest: Aerodrome AMO Strategy (Base)", async function () {
 
   describe("Harvest rewards", function () {
     it("Should be able to collect reward tokens", async () => {
-      const strategistAddr = await strategist.getAddress();
-
       await setERC20TokenBalance(
         aerodromeAmoStrategy.address,
         aero,
         "1337",
         hre
       );
-      const aeroBalanceBefore = await aero.balanceOf(strategistAddr);
-      await aerodromeAmoStrategy.connect(strategist).collectRewardTokens();
+      const aeroBalanceBefore = await aero.balanceOf(harvester.address);
+      await aerodromeAmoStrategy.connect(harvesterSigner).collectRewardTokens();
 
-      const aeroBalancediff = (await aero.balanceOf(strategistAddr)).sub(
+      const aeroBalancediff = (await aero.balanceOf(harvester.address)).sub(
         aeroBalanceBefore
       );
 

contracts/utils/addresses.js

@@ -282,6 +282,7 @@ addresses.arbitrumOne.WOETHProxy = "0xD8724322f44E5c58D7A815F542036fb17DbbF839";
 
 // Base
 addresses.base = {};
+addresses.base.HarvesterProxy = "0x247872f58f2fF11f9E8f89C1C48e460CfF0c6b29";
 addresses.base.BridgedWOETH = "0xD8724322f44E5c58D7A815F542036fb17DbbF839";
 addresses.base.AERO = "0x940181a94A35A4569E4529A3CDfB74e38FD98631";
 addresses.base.aeroRouterAddress = "0xcF77a3Ba9A5CA399B7c97c74d54e5b1Beb874E43";