fix: update tests

.forge-snapshots/FullRangeAddInitialLiquidity.snap

@@ -1 +1 @@
-412696
+390028

.forge-snapshots/FullRangeAddLiquidity.snap

@@ -1 +1 @@
-206962
+184294

.forge-snapshots/FullRangeFirstSwap.snap

@@ -1 +1 @@
-154763
+133445

.forge-snapshots/FullRangeInitialize.snap

@@ -1 +1 @@
-879542
+896690

.forge-snapshots/FullRangeRemoveLiquidity.snap

@@ -1 +1 @@
-200095
+177393

.forge-snapshots/FullRangeRemoveLiquidityAndRebalance.snap

@@ -1 +1 @@
-379287
+362063

.forge-snapshots/FullRangeSecondSwap.snap

@@ -1 +1 @@
-112303
+93791

.forge-snapshots/FullRangeSwap.snap

@@ -1 +1 @@
-153038
+131720

test/FullRange.t.sol

@@ -19,6 +19,7 @@ import {IHooks} from "@uniswap/v4-core/src/interfaces/IHooks.sol";
 import {UniswapV4ERC20} from "../contracts/libraries/UniswapV4ERC20.sol";
 import {FullMath} from "@uniswap/v4-core/src/libraries/FullMath.sol";
 import {SafeCast} from "@uniswap/v4-core/src/libraries/SafeCast.sol";
+import {HookEnabledSwapRouter} from "./utils/HookEnabledSwapRouter.sol";
 
 contract TestFullRange is Test, Deployers, GasSnapshot {
     using PoolIdLibrary for PoolKey;
@@ -47,6 +48,7 @@ contract TestFullRange is Test, Deployers, GasSnapshot {
         uint24 fee
     );
 
+    HookEnabledSwapRouter router;
     /// @dev Min tick for full range with tick spacing of 60
     int24 internal constant MIN_TICK = -887220;
     /// @dev Max tick for full range with tick spacing of 60
@@ -77,15 +79,20 @@ contract TestFullRange is Test, Deployers, GasSnapshot {
 
     function setUp() public {
         deployFreshManagerAndRouters();
+        router = new HookEnabledSwapRouter(manager);
         (currency0, currency1) = deployMintAndApprove2Currencies();
 
         token0 = new MockERC20("TestA", "A", 18);
+        token0.mint(address(this), 2 ** 128);
         token1 = new MockERC20("TestB", "B", 18);
+        token1.mint(address(this), 2 ** 128);
         token2 = new MockERC20("TestC", "C", 18);
+        token2.mint(address(this), 2 ** 128);
 
         FullRangeImplementation impl = new FullRangeImplementation(manager, fullRange);
         vm.etch(address(fullRange), address(impl).code);
 
+        key = createPoolKey(token0, token1);
         id = key.toId();
 
         key2 = createPoolKey(token1, token2);
@@ -97,11 +104,13 @@ contract TestFullRange is Test, Deployers, GasSnapshot {
         token0.approve(address(fullRange), type(uint256).max);
         token1.approve(address(fullRange), type(uint256).max);
         token2.approve(address(fullRange), type(uint256).max);
-        token0.approve(address(swapRouter), type(uint256).max);
-        token1.approve(address(swapRouter), type(uint256).max);
-        token2.approve(address(swapRouter), type(uint256).max);
+        token0.approve(address(router), type(uint256).max);
+        token1.approve(address(router), type(uint256).max);
+        token2.approve(address(router), type(uint256).max);
 
-        manager.initialize(keyWithLiq, SQRT_RATIO_1_1, ZERO_BYTES);
+        initPool(
+            keyWithLiq.currency0, keyWithLiq.currency1, fullRange, 3000, SQRT_RATIO_1_1, ZERO_BYTES
+        );
         fullRange.addLiquidity(
             FullRange.AddLiquidityParams(
                 keyWithLiq.currency0,
@@ -262,16 +271,16 @@ contract TestFullRange is Test, Deployers, GasSnapshot {
 
         vm.expectEmit(true, true, true, true);
         emit Swap(
-            id, address(swapRouter), 1 ether, -906610893880149131, 72045250990510446115798809072, 10 ether, -1901, 3000
+            id, address(router), 1 ether, -906610893880149131, 72045250990510446115798809072, 10 ether, -1901, 3000
         );
 
         IPoolManager.SwapParams memory params =
             IPoolManager.SwapParams({zeroForOne: true, amountSpecified: 1 ether, sqrtPriceLimitX96: SQRT_RATIO_1_2});
-        PoolSwapTest.TestSettings memory settings =
-            PoolSwapTest.TestSettings({withdrawTokens: true, settleUsingTransfer: true});
+        HookEnabledSwapRouter.TestSettings memory settings =
+            HookEnabledSwapRouter.TestSettings({withdrawTokens: true, settleUsingTransfer: true});
 
         snapStart("FullRangeSwap");
-        swapRouter.swap(key, params, settings, ZERO_BYTES);
+        router.swap(key, params, settings, ZERO_BYTES);
         snapEnd();
 
         (bool hasAccruedFees,) = fullRange.poolInfo(id);
@@ -305,10 +314,10 @@ contract TestFullRange is Test, Deployers, GasSnapshot {
 
         IPoolManager.SwapParams memory params =
             IPoolManager.SwapParams({zeroForOne: true, amountSpecified: 1000 ether, sqrtPriceLimitX96: SQRT_RATIO_1_2});
-        PoolSwapTest.TestSettings memory settings =
-            PoolSwapTest.TestSettings({withdrawTokens: true, settleUsingTransfer: true});
+        HookEnabledSwapRouter.TestSettings memory settings =
+            HookEnabledSwapRouter.TestSettings({withdrawTokens: true, settleUsingTransfer: true});
 
-        swapRouter.swap(key, params, settings, ZERO_BYTES);
+        router.swap(key, params, settings, ZERO_BYTES);
 
         vm.expectRevert(FullRange.TooMuchSlippage.selector);
         fullRange.addLiquidity(
@@ -330,18 +339,18 @@ contract TestFullRange is Test, Deployers, GasSnapshot {
 
         IPoolManager.SwapParams memory params =
             IPoolManager.SwapParams({zeroForOne: true, amountSpecified: 1 ether, sqrtPriceLimitX96: SQRT_RATIO_1_2});
-        PoolSwapTest.TestSettings memory settings =
-            PoolSwapTest.TestSettings({withdrawTokens: true, settleUsingTransfer: true});
+        HookEnabledSwapRouter.TestSettings memory settings =
+            HookEnabledSwapRouter.TestSettings({withdrawTokens: true, settleUsingTransfer: true});
 
         snapStart("FullRangeFirstSwap");
-        swapRouter.swap(testKey, params, settings, ZERO_BYTES);
+        router.swap(testKey, params, settings, ZERO_BYTES);
         snapEnd();
 
         (bool hasAccruedFees,) = fullRange.poolInfo(id);
         assertEq(hasAccruedFees, true);
 
         snapStart("FullRangeSecondSwap");
-        swapRouter.swap(testKey, params, settings, ZERO_BYTES);
+        router.swap(testKey, params, settings, ZERO_BYTES);
         snapEnd();
 
         (hasAccruedFees,) = fullRange.poolInfo(id);
@@ -366,11 +375,11 @@ contract TestFullRange is Test, Deployers, GasSnapshot {
         IPoolManager.SwapParams memory params =
             IPoolManager.SwapParams({zeroForOne: true, amountSpecified: 10000000, sqrtPriceLimitX96: SQRT_RATIO_1_2});
 
-        PoolSwapTest.TestSettings memory testSettings =
-            PoolSwapTest.TestSettings({withdrawTokens: true, settleUsingTransfer: true});
+        HookEnabledSwapRouter.TestSettings memory testSettings =
+            HookEnabledSwapRouter.TestSettings({withdrawTokens: true, settleUsingTransfer: true});
 
-        swapRouter.swap(key, params, testSettings, ZERO_BYTES);
-        swapRouter.swap(key2, params, testSettings, ZERO_BYTES);
+        router.swap(key, params, testSettings, ZERO_BYTES);
+        router.swap(key2, params, testSettings, ZERO_BYTES);
 
         (bool hasAccruedFees,) = fullRange.poolInfo(id);
         assertEq(hasAccruedFees, true);
@@ -549,10 +558,10 @@ contract TestFullRange is Test, Deployers, GasSnapshot {
         IPoolManager.SwapParams memory params =
             IPoolManager.SwapParams({zeroForOne: true, amountSpecified: 1 ether, sqrtPriceLimitX96: SQRT_RATIO_1_2});
 
-        PoolSwapTest.TestSettings memory testSettings =
-            PoolSwapTest.TestSettings({withdrawTokens: true, settleUsingTransfer: true});
+        HookEnabledSwapRouter.TestSettings memory testSettings =
+            HookEnabledSwapRouter.TestSettings({withdrawTokens: true, settleUsingTransfer: true});
 
-        swapRouter.swap(keyWithLiq, params, testSettings, ZERO_BYTES);
+        router.swap(keyWithLiq, params, testSettings, ZERO_BYTES);
 
         UniswapV4ERC20(liquidityToken).approve(address(fullRange), type(uint256).max);
 
@@ -676,10 +685,10 @@ contract TestFullRange is Test, Deployers, GasSnapshot {
         IPoolManager.SwapParams memory params =
             IPoolManager.SwapParams({zeroForOne: true, amountSpecified: 100 ether, sqrtPriceLimitX96: SQRT_RATIO_1_4});
 
-        PoolSwapTest.TestSettings memory testSettings =
-            PoolSwapTest.TestSettings({withdrawTokens: true, settleUsingTransfer: true});
+        HookEnabledSwapRouter.TestSettings memory testSettings =
+            HookEnabledSwapRouter.TestSettings({withdrawTokens: true, settleUsingTransfer: true});
 
-        swapRouter.swap(key, params, testSettings, ZERO_BYTES);
+        router.swap(key, params, testSettings, ZERO_BYTES);
 
         (bool hasAccruedFees,) = fullRange.poolInfo(id);
         assertEq(hasAccruedFees, true);
@@ -731,10 +740,10 @@ contract TestFullRange is Test, Deployers, GasSnapshot {
                 sqrtPriceLimitX96: SQRT_RATIO_1_4
             });
 
-            PoolSwapTest.TestSettings memory testSettings =
-                PoolSwapTest.TestSettings({withdrawTokens: true, settleUsingTransfer: true});
+            HookEnabledSwapRouter.TestSettings memory testSettings =
+                HookEnabledSwapRouter.TestSettings({withdrawTokens: true, settleUsingTransfer: true});
 
-            swapRouter.swap(key, params, testSettings, ZERO_BYTES);
+            router.swap(key, params, testSettings, ZERO_BYTES);
 
             // Test contract removes liquidity, succeeds
             UniswapV4ERC20(liquidityToken).approve(address(fullRange), type(uint256).max);

test/LimitOrder.t.sol

@@ -12,15 +12,16 @@ import {Deployers} from "@uniswap/v4-core/test/utils/Deployers.sol";
 import {TestERC20} from "@uniswap/v4-core/src/test/TestERC20.sol";
 import {CurrencyLibrary, Currency} from "@uniswap/v4-core/src/types/Currency.sol";
 import {PoolId, PoolIdLibrary} from "@uniswap/v4-core/src/types/PoolId.sol";
-import {PoolSwapTest} from "@uniswap/v4-core/src/test/PoolSwapTest.sol";
 import {TickMath} from "@uniswap/v4-core/src/libraries/TickMath.sol";
 import {PoolKey} from "@uniswap/v4-core/src/types/PoolKey.sol";
+import {HookEnabledSwapRouter} from "./utils/HookEnabledSwapRouter.sol";
 
 contract TestLimitOrder is Test, Deployers {
     using PoolIdLibrary for PoolKey;
 
     uint160 constant SQRT_RATIO_10_1 = 250541448375047931186413801569;
 
+    HookEnabledSwapRouter router;
     TestERC20 token0;
     TestERC20 token1;
     LimitOrder limitOrder = LimitOrder(address(uint160(Hooks.AFTER_INITIALIZE_FLAG | Hooks.AFTER_SWAP_FLAG)));
@@ -30,6 +31,7 @@ contract TestLimitOrder is Test, Deployers {
         deployFreshManagerAndRouters();
         (currency0, currency1) = deployMintAndApprove2Currencies();
 
+        router = new HookEnabledSwapRouter(manager);
         token0 = TestERC20(Currency.unwrap(currency0));
         token1 = TestERC20(Currency.unwrap(currency1));
 
@@ -45,14 +47,15 @@ contract TestLimitOrder is Test, Deployers {
             }
         }
 
-        key = PoolKey(currency0, currency1, 3000, 60, limitOrder);
-        id = key.toId();
-        manager.initialize(key, SQRT_RATIO_1_1, ZERO_BYTES);
+        // key = PoolKey(currency0, currency1, 3000, 60, limitOrder);
+        (key, id) = initPoolAndAddLiquidity(
+            currency0, currency1, limitOrder, 3000, SQRT_RATIO_1_1, ZERO_BYTES
+        );
 
         token0.approve(address(limitOrder), type(uint256).max);
         token1.approve(address(limitOrder), type(uint256).max);
-        token0.approve(address(swapRouter), type(uint256).max);
-        token1.approve(address(swapRouter), type(uint256).max);
+        token0.approve(address(router), type(uint256).max);
+        token1.approve(address(router), type(uint256).max);
     }
 
     function testGetTickLowerLast() public {
@@ -100,10 +103,10 @@ contract TestLimitOrder is Test, Deployers {
 
     function testZeroForOneInRangeRevert() public {
         // swapping is free, there's no liquidity in the pool, so we only need to specify 1 wei
-        swapRouter.swap(
+        router.swap(
             key,
-            IPoolManager.SwapParams(false, 1, SQRT_RATIO_1_1 + 1),
-            PoolSwapTest.TestSettings(true, true),
+            IPoolManager.SwapParams(false, 1 ether, SQRT_RATIO_1_1 + 1),
+            HookEnabledSwapRouter.TestSettings(true, true),
             ZERO_BYTES
         );
         vm.expectRevert(LimitOrder.InRange.selector);
@@ -126,8 +129,8 @@ contract TestLimitOrder is Test, Deployers {
 
     function testNotZeroForOneInRangeRevert() public {
         // swapping is free, there's no liquidity in the pool, so we only need to specify 1 wei
-        swapRouter.swap(
-            key, IPoolManager.SwapParams(true, 1, SQRT_RATIO_1_1 - 1), PoolSwapTest.TestSettings(true, true), ZERO_BYTES
+        router.swap(
+            key, IPoolManager.SwapParams(true, 1, SQRT_RATIO_1_1 - 1), HookEnabledSwapRouter.TestSettings(true, true), ZERO_BYTES
         );
         vm.expectRevert(LimitOrder.InRange.selector);
         limitOrder.place(key, -60, false, 1000000);
@@ -185,10 +188,10 @@ contract TestLimitOrder is Test, Deployers {
         uint128 liquidity = 1000000;
         limitOrder.place(key, tickLower, zeroForOne, liquidity);
 
-        swapRouter.swap(
+        router.swap(
             key,
             IPoolManager.SwapParams(false, 1e18, TickMath.getSqrtRatioAtTick(60)),
-            PoolSwapTest.TestSettings(true, true),
+            HookEnabledSwapRouter.TestSettings(true, true),
             ZERO_BYTES
         );
 

test/TWAMM.t.sol

@@ -93,6 +93,7 @@ contract TWAMMTest is Test, Deployers, GasSnapshot {
         (PoolKey memory initKey, PoolId initId) = newPoolKeyWithTWAMM(twamm);
         assertEq(twamm.lastVirtualOrderTimestamp(initId), 0);
         vm.warp(10000);
+
         manager.initialize(initKey, SQRT_RATIO_1_1, ZERO_BYTES);
         assertEq(twamm.lastVirtualOrderTimestamp(initId), 10000);
     }
@@ -185,8 +186,8 @@ contract TWAMMTest is Test, Deployers, GasSnapshot {
 
         assertEq(sellRate0For1, 2e18 / (expiration2 - submitTimestamp2));
         assertEq(sellRate1For0, 3e18 / (expiration2 - submitTimestamp1));
-        assertEq(earningsFactor0For1, 1712020976636017581269515821040000);
-        assertEq(earningsFactor1For0, 1470157410324350030712806974476955);
+        assertEq(earningsFactor0For1, 1636776489931663248324424309240000);
+        assertEq(earningsFactor1For0, 1534530274609724617872321172427618);
     }
 
     function testTWAMM_submitOrder_EmitsEvent() public {
@@ -409,8 +410,8 @@ contract TWAMMTest is Test, Deployers, GasSnapshot {
     }
 
     function newPoolKeyWithTWAMM(IHooks hooks) public returns (PoolKey memory, PoolId) {
-        MockERC20[] memory tokens = deployTokens(2, 2 ** 255);
-        PoolKey memory key = PoolKey(Currency.wrap(address(tokens[0])), Currency.wrap(address(tokens[1])), 0, 60, hooks);
+        (Currency _token0, Currency _token1) = deployMintAndApprove2Currencies();
+        PoolKey memory key = PoolKey(_token0, _token1, 0, 60, hooks);
         return (key, key.toId());
     }
 

test/utils/HookEnabledSwapRouter.sol

@@ -0,0 +1,65 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.20;
+
+import {CurrencyLibrary, Currency} from "@uniswap/v4-core/src/types/Currency.sol";
+import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
+import {BalanceDelta} from "@uniswap/v4-core/src/types/BalanceDelta.sol";
+import {PoolKey} from "@uniswap/v4-core/src/types/PoolKey.sol";
+import {PoolTestBase} from "@uniswap/v4-core/src/test/PoolTestBase.sol";
+import {Test} from "forge-std/Test.sol";
+
+contract HookEnabledSwapRouter is Test, PoolTestBase {
+    using CurrencyLibrary for Currency;
+
+    constructor(IPoolManager _manager) PoolTestBase(_manager) {}
+
+    error NoSwapOccurred();
+
+    struct CallbackData {
+        address sender;
+        TestSettings testSettings;
+        PoolKey key;
+        IPoolManager.SwapParams params;
+        bytes hookData;
+    }
+
+    struct TestSettings {
+        bool withdrawTokens;
+        bool settleUsingTransfer;
+    }
+
+    function swap(
+        PoolKey memory key,
+        IPoolManager.SwapParams memory params,
+        TestSettings memory testSettings,
+        bytes memory hookData
+    ) external payable returns (BalanceDelta delta) {
+        delta = abi.decode(
+            manager.lock(abi.encode(CallbackData(msg.sender, testSettings, key, params, hookData))), (BalanceDelta)
+        );
+
+        uint256 ethBalance = address(this).balance;
+        if (ethBalance > 0) CurrencyLibrary.NATIVE.transfer(msg.sender, ethBalance);
+    }
+
+    function lockAcquired(bytes calldata rawData) external returns (bytes memory) {
+        require(msg.sender == address(manager));
+
+        CallbackData memory data = abi.decode(rawData, (CallbackData));
+
+        BalanceDelta delta = manager.swap(data.key, data.params, data.hookData);
+
+        // Make sure youve added liquidity to the test pool!
+        if (BalanceDelta.unwrap(delta) == 0) revert NoSwapOccurred();
+
+        if (data.params.zeroForOne) {
+            _settle(data.key.currency0, data.sender, delta.amount0(), data.testSettings.settleUsingTransfer);
+            _take(data.key.currency1, data.sender, delta.amount1(), data.testSettings.withdrawTokens);
+        } else {
+            _settle(data.key.currency1, data.sender, delta.amount1(), data.testSettings.settleUsingTransfer);
+            _take(data.key.currency0, data.sender, delta.amount0(), data.testSettings.withdrawTokens);
+        }
+
+        return abi.encode(delta);
+    }
+}