Angstrom review / refactor

contracts/test/AngstromBalancerMock.sol

@@ -19,16 +19,8 @@ contract AngstromBalancerMock is AngstromBalancer {
         // solhint-disable-previous-line no-empty-blocks
     }
 
-    function manualUnlockAngstromWithRouter() external {
-        _unlockAngstromWithRouter();
-    }
-
-    function getLastUnlockBlockNumber() external view returns (uint256) {
-        return _lastUnlockBlockNumber;
-    }
-
-    function isNode(address account) external view returns (bool) {
-        return _isNode(account);
+    function manualUnlockAngstrom() external {
+        _unlockAngstrom();
     }
 
     function getDigest() external view returns (bytes32) {

test/foundry/AngstromBalancerUnit.t.sol

@@ -7,78 +7,51 @@ import "forge-std/Test.sol";
 import { IAuthentication } from "@balancer-labs/v3-interfaces/contracts/solidity-utils/helpers/IAuthentication.sol";
 
 import { ArrayHelpers } from "@balancer-labs/v3-solidity-utils/contracts/test/ArrayHelpers.sol";
-import { BaseVaultTest } from "@balancer-labs/v3-vault/test/foundry/utils/BaseVaultTest.sol";
 
 import { AngstromBalancerMock } from "../../contracts/test/AngstromBalancerMock.sol";
 import { AngstromBalancer } from "../../contracts/AngstromBalancer.sol";
+import { BaseAngstromTest } from "./utils/BaseAngstromTest.sol";
 
-contract AngstromBalancerUnitTest is BaseVaultTest {
+contract AngstromBalancerUnitTest is BaseAngstromTest {
     using ArrayHelpers for *;
 
-    AngstromBalancerMock private _angstromBalancer;
-
-    function setUp() public virtual override {
-        super.setUp();
-
-        _angstromBalancer = new AngstromBalancerMock(vault, weth, permit2, "AngstromBalancer Mock v1");
-
-        authorizer.grantRole(_angstromBalancer.getActionId(AngstromBalancer.toggleNodes.selector), admin);
-    }
-
     function testToggleNodesIsAuthenticated() public {
         vm.expectRevert(IAuthentication.SenderNotAllowed.selector);
-        _angstromBalancer.toggleNodes([bob].toMemoryArray());
+        angstromBalancer.toggleNodes([bob].toMemoryArray());
     }
 
     function testToggleNodes() public {
-        vm.prank(admin);
-        _angstromBalancer.toggleNodes([bob].toMemoryArray());
-        assertTrue(_angstromBalancer.isNode(bob), "Bob is not a node");
+        makeAngstromNode(bob);
+        assertTrue(angstromBalancer.isRegisteredNode(bob), "Bob is not a node");
     }
 
     function testAddAndRemoveNodes() public {
-        vm.startPrank(admin);
-        _angstromBalancer.toggleNodes([bob, alice].toMemoryArray());
-        assertTrue(_angstromBalancer.isNode(bob), "Bob is not a node");
-        assertTrue(_angstromBalancer.isNode(alice), "Alice is not a node");
-        _angstromBalancer.toggleNodes([bob].toMemoryArray());
-        vm.stopPrank();
-        assertFalse(_angstromBalancer.isNode(bob), "Bob is still a node");
-        assertTrue(_angstromBalancer.isNode(alice), "Alice is not a node after bob was removed");
-    }
+        makeAngstromNodes([bob, alice].toMemoryArray());
+        assertTrue(angstromBalancer.isRegisteredNode(bob), "Bob is not a node");
+        assertTrue(angstromBalancer.isRegisteredNode(alice), "Alice is not a node");
 
-    function testUnlockAngstromNotNode() public {
-        vm.expectRevert(AngstromBalancer.NotNode.selector);
-        _angstromBalancer.manualUnlockAngstromWithRouter();
-    }
-
-    function testUnlockAngstromTwice() public {
         vm.prank(admin);
-        _angstromBalancer.toggleNodes([bob].toMemoryArray());
+        angstromBalancer.toggleNodes([bob].toMemoryArray());
 
-        vm.startPrank(bob);
-        _angstromBalancer.manualUnlockAngstromWithRouter();
-        vm.expectRevert(AngstromBalancer.OnlyOncePerBlock.selector);
-        _angstromBalancer.manualUnlockAngstromWithRouter();
-        vm.stopPrank();
+        assertFalse(angstromBalancer.isRegisteredNode(bob), "Bob is still a node");
+        assertTrue(angstromBalancer.isRegisteredNode(alice), "Alice is not a node after bob was removed");
     }
 
     function testUnlockAngstromSetsLastUnlockBlockNumber() public {
-        vm.prank(admin);
-        _angstromBalancer.toggleNodes([bob].toMemoryArray());
+        makeAngstromNode(bob);
 
-        assertEq(_angstromBalancer.getLastUnlockBlockNumber(), 0, "Last unlock block number is not 0");
+        assertEq(angstromBalancer.getLastUnlockBlockNumber(), 0, "Last unlock block number is not 0");
 
         vm.prank(bob);
-        _angstromBalancer.manualUnlockAngstromWithRouter();
+        angstromBalancer.manualUnlockAngstrom();
         assertEq(
-            _angstromBalancer.getLastUnlockBlockNumber(),
+            angstromBalancer.getLastUnlockBlockNumber(),
             block.number,
             "Last unlock block number is not the current block number"
         );
     }
 
     function testGetVault() public view {
-        assertEq(address(_angstromBalancer.getVault()), address(vault), "Wrong vault address");
+        assertEq(address(angstromBalancer.getVault()), address(vault), "Wrong vault address");
     }
 }

test/foundry/AngstromHook.t.sol

@@ -13,38 +13,23 @@ import {
 } from "@balancer-labs/v3-interfaces/contracts/vault/BatchRouterTypes.sol";
 
 import { ArrayHelpers } from "@balancer-labs/v3-solidity-utils/contracts/test/ArrayHelpers.sol";
-import { BaseVaultTest } from "@balancer-labs/v3-vault/test/foundry/utils/BaseVaultTest.sol";
 import { FixedPoint } from "@balancer-labs/v3-solidity-utils/contracts/math/FixedPoint.sol";
 
 import { AngstromBalancerMock } from "../../contracts/test/AngstromBalancerMock.sol";
 import { AngstromBalancer } from "../../contracts/AngstromBalancer.sol";
+import { BaseAngstromTest } from "./utils/BaseAngstromTest.sol";
 
-contract AngstromHookTest is BaseVaultTest {
+contract AngstromHookTest is BaseAngstromTest {
     using ArrayHelpers for *;
 
-    AngstromBalancerMock private _angstromBalancer;
-
-    function setUp() public virtual override {
-        super.setUp();
-        authorizer.grantRole(_angstromBalancer.getActionId(AngstromBalancer.toggleNodes.selector), admin);
-    }
-
-    function createHook() internal override returns (address) {
-        // Creating the router and hook in this function ensures that "pool" has the hook set correctly.
-        _angstromBalancer = new AngstromBalancerMock(vault, weth, permit2, "AngstromBalancer Mock v1");
-        return address(_angstromBalancer);
-    }
-
     /***************************************************************************
                                     Swaps
     ***************************************************************************/
 
     function testOnBeforeSwapNotNode() public {
-        (, bytes memory userData) = _generateSignatureAndUserData(bob, bobKey);
-
         vm.expectRevert(AngstromBalancer.NotNode.selector);
         vm.prank(bob);
-        router.swapSingleTokenExactIn(pool, dai, usdc, 1e18, 0, MAX_UINT256, false, userData);
+        router.swapSingleTokenExactIn(pool, dai, usdc, 1e18, 0, MAX_UINT256, false, bobUserData);
     }
 
     function testOnBeforeSwapCannotSwapWhileLocked() public {
@@ -64,55 +49,37 @@ contract AngstromHookTest is BaseVaultTest {
 
     function testOnBeforeSwapInvalidSignature() public {
         // If the signature is invalid, the hook reverts (in this case, signer and key do not match).
-        (, bytes memory userData) = _generateSignatureAndUserData(bob, aliceKey);
+        (, bytes memory userData) = generateSignatureAndUserData(bob, aliceKey);
 
         vm.prank(admin);
-        _angstromBalancer.toggleNodes([bob].toMemoryArray());
+        angstromBalancer.toggleNodes([bob].toMemoryArray());
 
         vm.expectRevert(AngstromBalancer.InvalidSignature.selector);
         vm.prank(bob);
         router.swapSingleTokenExactIn(pool, dai, usdc, 1e18, 0, MAX_UINT256, false, userData);
     }
 
     function testOnBeforeSwapSucceedsAndSetBlockNumber() public {
-        (, bytes memory userData) = _generateSignatureAndUserData(bob, bobKey);
-
-        vm.prank(admin);
-        _angstromBalancer.toggleNodes([bob].toMemoryArray());
+        makeAngstromNode(bob);
 
         vm.prank(bob);
-        router.swapSingleTokenExactIn(pool, dai, usdc, 1e18, 0, MAX_UINT256, false, userData);
+        router.swapSingleTokenExactIn(pool, dai, usdc, 1e18, 0, MAX_UINT256, false, bobUserData);
         assertEq(
-            _angstromBalancer.getLastUnlockBlockNumber(),
+            angstromBalancer.getLastUnlockBlockNumber(),
             block.number,
             "Last unlock block number is not the current block number"
         );
     }
 
     function testOnlyOncePerBlock() public {
-        vm.prank(admin);
-        _angstromBalancer.toggleNodes([bob].toMemoryArray());
-
-        (bytes memory signature, ) = _generateSignatureAndUserData(bob, bobKey);
-        vm.prank(bob);
-        _angstromBalancer.unlockWithEmptyAttestation(bob, signature);
-
-        vm.expectRevert(AngstromBalancer.OnlyOncePerBlock.selector);
-        vm.prank(bob);
-        _angstromBalancer.unlockWithEmptyAttestation(bob, signature);
-    }
-
-    function testOnlyOncePerBlockCalldata() public {
-        vm.prank(admin);
-        _angstromBalancer.toggleNodes([bob].toMemoryArray());
+        makeAngstromNode(bob);
 
-        (bytes memory signature, ) = _generateSignatureAndUserData(bob, bobKey);
         vm.prank(bob);
-        _angstromBalancer.unlockWithEmptyAttestationCalldata(bob, signature);
+        angstromBalancer.unlockWithEmptyAttestation(bob, bobSignature);
 
         vm.expectRevert(AngstromBalancer.OnlyOncePerBlock.selector);
         vm.prank(bob);
-        _angstromBalancer.unlockWithEmptyAttestationCalldata(bob, signature);
+        angstromBalancer.unlockWithEmptyAttestation(bob, bobSignature);
     }
 
     /***************************************************************************
@@ -138,40 +105,40 @@ contract AngstromHookTest is BaseVaultTest {
     }
 
     function testOnBeforeAddLiquidityUnbalancedNotNode() public {
-        (, bytes memory userData) = _generateSignatureAndUserData(alice, aliceKey);
-
         vm.expectRevert(AngstromBalancer.NotNode.selector);
         vm.prank(alice);
-        router.addLiquidityUnbalanced(pool, [FixedPoint.ONE, FixedPoint.ONE].toMemoryArray(), 1e18, false, userData);
+        router.addLiquidityUnbalanced(
+            pool,
+            [FixedPoint.ONE, FixedPoint.ONE].toMemoryArray(),
+            1e18,
+            false,
+            aliceUserData
+        );
     }
 
     function testOnBeforeAddLiquidityUnbalancedInvalidSignature() public {
-        (, bytes memory userData) = _generateSignatureAndUserData(alice, bobKey);
+        (, bytes memory userData) = generateSignatureAndUserData(alice, bobKey);
 
-        vm.prank(admin);
-        _angstromBalancer.toggleNodes([alice].toMemoryArray());
+        makeAngstromNode(alice);
 
         vm.expectRevert(AngstromBalancer.InvalidSignature.selector);
         vm.prank(alice);
         router.addLiquidityUnbalanced(pool, [FixedPoint.ONE, FixedPoint.ONE].toMemoryArray(), 1e18, false, userData);
     }
 
     function testOnBeforeAddLiquidityUnbalancedSucceedsAndSetBlockNumber() public {
-        (, bytes memory userData) = _generateSignatureAndUserData(alice, aliceKey);
-
-        vm.prank(admin);
-        _angstromBalancer.toggleNodes([alice].toMemoryArray());
+        makeAngstromNode(alice);
 
         vm.prank(alice);
         uint256 expectedBptAmountOut = router.addLiquidityUnbalanced(
             pool,
             [FixedPoint.ONE, FixedPoint.ONE].toMemoryArray(),
             1e18,
             false,
-            userData
+            aliceUserData
         );
         assertEq(
-            _angstromBalancer.getLastUnlockBlockNumber(),
+            angstromBalancer.getLastUnlockBlockNumber(),
             block.number,
             "Last unlock block number is not the current block number"
         );
@@ -206,52 +173,36 @@ contract AngstromHookTest is BaseVaultTest {
     }
 
     function testOnBeforeRemoveLiquidityUnbalancedNotNode() public {
-        (, bytes memory userData) = _generateSignatureAndUserData(lp, lpKey);
-
         vm.expectRevert(AngstromBalancer.NotNode.selector);
         vm.prank(lp);
-        router.removeLiquiditySingleTokenExactIn(pool, 1e18, dai, 0.1e18, false, userData);
+        router.removeLiquiditySingleTokenExactIn(pool, 1e18, dai, 0.1e18, false, lpUserData);
     }
 
     function testOnBeforeRemoveLiquidityUnbalancedInvalidSignature() public {
-        (, bytes memory userData) = _generateSignatureAndUserData(lp, bobKey);
+        (, bytes memory userData) = generateSignatureAndUserData(lp, bobKey);
 
-        vm.prank(admin);
-        _angstromBalancer.toggleNodes([lp].toMemoryArray());
+        makeAngstromNode(lp);
 
         vm.expectRevert(AngstromBalancer.InvalidSignature.selector);
         vm.prank(lp);
         router.removeLiquiditySingleTokenExactIn(pool, 1e18, dai, 0.1e18, false, userData);
     }
 
     function testOnBeforeRemoveLiquidityUnbalancedSucceedsAndSetBlockNumber() public {
-        (, bytes memory userData) = _generateSignatureAndUserData(lp, lpKey);
-
-        vm.prank(admin);
-        _angstromBalancer.toggleNodes([lp].toMemoryArray());
+        makeAngstromNode(lp);
 
         Balances memory balancesBefore = getBalances(lp);
 
         vm.prank(lp);
-        router.removeLiquiditySingleTokenExactIn(pool, 1e18, dai, 0.1e18, false, userData);
+        router.removeLiquiditySingleTokenExactIn(pool, 1e18, dai, 0.1e18, false, lpUserData);
 
         Balances memory balancesAfter = getBalances(lp);
 
         assertEq(
-            _angstromBalancer.getLastUnlockBlockNumber(),
+            angstromBalancer.getLastUnlockBlockNumber(),
             block.number,
             "Last unlock block number is not the current block number"
         );
         assertEq(balancesAfter.lpBpt, balancesBefore.lpBpt - 1e18, "LP did not burn BPTs");
     }
-
-    function _generateSignatureAndUserData(
-        address signer,
-        uint256 privateKey
-    ) private view returns (bytes memory signature, bytes memory userData) {
-        bytes32 hash = _angstromBalancer.getDigest();
-        (uint8 v, bytes32 r, bytes32 s) = vm.sign(privateKey, hash);
-        signature = abi.encodePacked(r, s, v);
-        userData = abi.encodePacked(signer, signature);
-    }
 }