update constructor to take owner

Uniswap · Jun 13, 2024 · 0a61b06 · 0a61b06
1 parent da0ce08
commit 0a61b06
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 6 deletions.
diff --git a/contracts/hooks/examples/FeeTaking.sol b/contracts/hooks/examples/FeeTaking.sol
@@ -24,7 +24,7 @@ contract FeeTaking is BaseHook, IUnlockCallback, Owned {
         Currency[] currencies;
     }
 
-    constructor(IPoolManager _poolManager, uint128 _swapFeeBips) BaseHook(_poolManager) Owned(msg.sender) {
+    constructor(IPoolManager _poolManager, uint128 _swapFeeBips, address _owner) BaseHook(_poolManager) Owned(_owner) {
         swapFeeBips = _swapFeeBips;
     }
 

diff --git a/test/FeeTaking.t.sol b/test/FeeTaking.t.sol
@@ -43,7 +43,7 @@ contract FeeTakingTest is Test, Deployers {
         token1 = TestERC20(Currency.unwrap(currency1));
 
         vm.record();
-        FeeTakingImplementation impl = new FeeTakingImplementation(manager, 25, feeTaking);
+        FeeTakingImplementation impl = new FeeTakingImplementation(manager, 25, address(this), feeTaking);
         (, bytes32[] memory writes) = vm.accesses(address(impl));
         vm.etch(address(feeTaking), address(impl).code);
         // for each storage key that was written during the hook implementation, copy the value over
@@ -100,19 +100,23 @@ contract FeeTakingTest is Test, Deployers {
         currencies[0] = key.currency0;
         currencies[1] = key.currency1;
         feeTaking.withdraw(TREASURY, currencies);
-        assertEq(manager.balanceOf(address(this), CurrencyLibrary.toId(key.currency0)), 0);
-        assertEq(manager.balanceOf(address(this), CurrencyLibrary.toId(key.currency1)), 0);
+        assertEq(manager.balanceOf(address(feeTaking), CurrencyLibrary.toId(key.currency0)), 0);
+        assertEq(manager.balanceOf(address(feeTaking), CurrencyLibrary.toId(key.currency1)), 0);
         assertEq(currency0.balanceOf(TREASURY) / R, expectedFee2 / R);
         assertEq(currency1.balanceOf(TREASURY) / R, expectedFee / R);
     }
 
+    // this would error had the hook not used ERC6909
     function testEdgeCase() public {
         // first, deplete the pool of token1
         // Swap exact token0 for token1 //
         bool zeroForOne = true;
         int256 amountSpecified = -1e18;
         BalanceDelta swapDelta = swap(key, zeroForOne, amountSpecified, ZERO_BYTES);
         // ---------------------------- //
+        // now, pool only has 1 wei of token1
+        uint256 poolToken1 = currency1.balanceOf(address(manager)) - manager.balanceOf(address(feeTaking), CurrencyLibrary.toId(key.currency1));
+        assertEq(poolToken1, 1);
 
         uint128 output = uint128(swapDelta.amount1());
         assertTrue(output > 0);

diff --git a/test/shared/implementation/FeeTakingImplementation.sol b/test/shared/implementation/FeeTakingImplementation.sol
@@ -7,8 +7,8 @@ import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
 import {Hooks} from "@uniswap/v4-core/src/libraries/Hooks.sol";
 
 contract FeeTakingImplementation is FeeTaking {
-    constructor(IPoolManager _poolManager, uint128 _swapFeeBips, FeeTaking addressToEtch)
-        FeeTaking(_poolManager, _swapFeeBips)
+    constructor(IPoolManager _poolManager, uint128 _swapFeeBips, address _owner, FeeTaking addressToEtch)
+        FeeTaking(_poolManager, _swapFeeBips, _owner)
     {
         Hooks.validateHookPermissions(addressToEtch, getHookPermissions());
     }