Unit Tests: Functions in the Swap path

pkg/interfaces/contracts/test/IVaultMainMock.sol

@@ -25,6 +25,32 @@ interface IVaultMainMock {
         address pauseManager
     ) external;
 
+    function manualSetLockers(address[] memory lockers) external;
+
+    function manualSetInitializedPool(address pool, bool isPoolInitialized) external;
+
+    function manualSetPoolPaused(address, bool, uint256) external;
+
+    function manualSetVaultState(bool, bool, uint256, uint256) external;
+
+    function manualSetPoolTokenConfig(address, IERC20[] memory, TokenConfig[] memory) external;
+
+    function manualSetPoolConfig(address, PoolConfig memory) external;
+
+    function manualSetPoolTokenBalances(address, IERC20[] memory, uint256[] memory) external;
+
+    function testWithLocker() external view;
+
+    function testWithInitializedPool(address pool) external view;
+
+    function testEnsurePoolNotPaused(address) external view;
+
+    function testEnsureUnpausedAndGetVaultState(address) external view returns (VaultState memory);
+
+    function testInternalGetPoolTokenInfo(
+        address
+    ) external view returns (TokenConfig[] memory, uint256[] memory, uint256[] memory, PoolConfig memory);
+
     function getDecimalScalingFactors(address pool) external view returns (uint256[] memory);
 
     function recoveryModeExit(address pool) external view;

pkg/vault/contracts/test/VaultMock.sol

@@ -28,6 +28,7 @@ contract VaultMock is IVaultMainMock, Vault {
     using ScalingHelpers for uint256;
     using PackedTokenBalance for bytes32;
     using PoolConfigLib for PoolConfig;
+    using VaultStateLib for VaultState;
 
     PoolFactoryMock private immutable _poolFactoryMock;
 
@@ -99,6 +100,85 @@ contract VaultMock is IVaultMainMock, Vault {
         );
     }
 
+    function manualSetLockers(address[] memory lockers) public {
+        _lockers = lockers;
+    }
+
+    function manualSetInitializedPool(address pool, bool isPoolInitialized) public {
+        PoolConfig memory poolConfig = _poolConfig[pool].toPoolConfig();
+        poolConfig.isPoolInitialized = isPoolInitialized;
+        _poolConfig[pool] = poolConfig.fromPoolConfig();
+    }
+
+    function manualSetPoolPaused(address pool, bool isPoolPaused, uint256 pauseWindowEndTime) public {
+        PoolConfig memory poolConfig = _poolConfig[pool].toPoolConfig();
+        poolConfig.isPoolPaused = isPoolPaused;
+        poolConfig.pauseWindowEndTime = pauseWindowEndTime;
+        _poolConfig[pool] = poolConfig.fromPoolConfig();
+    }
+
+    function manualSetVaultState(
+        bool isVaultPaused,
+        bool isQueryDisabled,
+        uint256 protocolSwapFeePercentage,
+        uint256 protocolYieldFeePercentage
+    ) public {
+        VaultState memory vaultState = _vaultState.toVaultState();
+        vaultState.isVaultPaused = isVaultPaused;
+        vaultState.isQueryDisabled = isQueryDisabled;
+        vaultState.protocolSwapFeePercentage = protocolSwapFeePercentage;
+        vaultState.protocolYieldFeePercentage = protocolYieldFeePercentage;
+        _vaultState = vaultState.fromVaultState();
+    }
+
+    function manualSetPoolConfig(address pool, PoolConfig memory poolConfig) public {
+        _poolConfig[pool] = poolConfig.fromPoolConfig();
+    }
+
+    // This function only works flawlessly for pools that are not registered by the vault,
+    // since mapping entries can't be deleted, only altered.
+    function manualSetPoolTokenConfig(address pool, IERC20[] memory tokens, TokenConfig[] memory tokenConfig) public {
+        for (uint256 i = 0; i < tokens.length; i++) {
+            _poolTokenConfig[pool][tokens[i]] = tokenConfig[i];
+        }
+    }
+
+    // This function only works for pools that are not registered by the vault,
+    // since mapping entries can't be deleted, only altered.
+    function manualSetPoolTokenBalances(address pool, IERC20[] memory tokens, uint256[] memory tokenBalanceRaw) public {
+        EnumerableMap.IERC20ToBytes32Map storage poolTokenBalances = _poolTokenBalances[pool];
+        for (uint256 i = 0; i < tokens.length; i++) {
+            poolTokenBalances.set(tokens[i], bytes32(tokenBalanceRaw[i]));
+        }
+    }
+
+    function testWithLocker() public view withLocker {}
+
+    function testWithInitializedPool(address pool) public view withInitializedPool(pool) {}
+
+    function testEnsurePoolNotPaused(address pool) public view {
+        _ensurePoolNotPaused(pool);
+    }
+
+    function testEnsureUnpausedAndGetVaultState(address pool) public view returns (VaultState memory vaultState) {
+        vaultState = _ensureUnpausedAndGetVaultState(pool);
+    }
+
+    function testInternalGetPoolTokenInfo(
+        address pool
+    )
+        public
+        view
+        returns (
+            TokenConfig[] memory tokenConfig,
+            uint256[] memory balancesRaw,
+            uint256[] memory decimalScalingFactors,
+            PoolConfig memory poolConfig
+        )
+    {
+        (tokenConfig, balancesRaw, decimalScalingFactors, poolConfig) = _getPoolTokenInfo(pool);
+    }
+
     function buildTokenConfig(IERC20[] memory tokens) public pure returns (TokenConfig[] memory tokenConfig) {
         tokenConfig = new TokenConfig[](tokens.length);
         for (uint256 i = 0; i < tokens.length; i++) {

pkg/vault/test/foundry/unit/PoolAndVaultPaused.t.sol

@@ -0,0 +1,76 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+pragma solidity ^0.8.4;
+
+import "@balancer-labs/v3-interfaces/contracts/vault/IVaultErrors.sol";
+import "@balancer-labs/v3-interfaces/contracts/vault/VaultTypes.sol";
+
+import { BaseVaultTest } from "../utils/BaseVaultTest.sol";
+
+contract PoolAndVaultPausedTest is BaseVaultTest {
+    // The pauseWindowEndTime is stored as a 32 bits number
+    uint256 private constant _MAX_PAUSE_END_TIME = 2 ** 32 - 1;
+
+    function setUp() public virtual override {
+        BaseVaultTest.setUp();
+    }
+
+    /*******************************************************************************
+                                _ensurePoolNotPaused
+    *******************************************************************************/
+
+    function testPausedPoolWithBigEndTime() public {
+        vault.manualSetPoolPaused(address(pool), true, _MAX_PAUSE_END_TIME);
+        vm.expectRevert(abi.encodeWithSelector(IVaultErrors.PoolPaused.selector, address(pool)));
+        vault.testEnsurePoolNotPaused(address(pool));
+    }
+
+    function testPausedPoolWithSmallEndTime() public {
+        vault.manualSetPoolPaused(address(pool), true, 1);
+        vault.testEnsurePoolNotPaused(address(pool));
+    }
+
+    function testUnpausedPoolWithBigEndTime() public {
+        vault.manualSetPoolPaused(address(pool), false, _MAX_PAUSE_END_TIME);
+        vault.testEnsurePoolNotPaused(address(pool));
+    }
+
+    function testUnpausedPoolWithSmallEndTime() public {
+        vault.manualSetPoolPaused(address(pool), false, 1);
+        vault.testEnsurePoolNotPaused(address(pool));
+    }
+
+    /*******************************************************************************
+                           _ensureUnpausedAndGetVaultState
+    *******************************************************************************/
+
+    function testVaultPaused() public {
+        vault.manualSetVaultState(true, false, 1e16, 1e16);
+        vm.expectRevert(abi.encodeWithSelector(IVaultErrors.VaultPaused.selector));
+        vault.testEnsureUnpausedAndGetVaultState(address(pool));
+    }
+
+    function testVaultUnpausedButPoolPaused() public {
+        vault.manualSetVaultState(false, false, 1e16, 1e16);
+        vault.manualSetPoolPaused(address(pool), true, _MAX_PAUSE_END_TIME);
+        vm.expectRevert(abi.encodeWithSelector(IVaultErrors.PoolPaused.selector, address(pool)));
+        vault.testEnsureUnpausedAndGetVaultState(address(pool));
+    }
+
+    function testVaultPausedButPoolUnpaused() public {
+        vault.manualSetVaultState(true, false, 1e16, 1e16);
+        vault.manualSetPoolPaused(address(pool), false, 1);
+        vm.expectRevert(abi.encodeWithSelector(IVaultErrors.VaultPaused.selector));
+        vault.testEnsureUnpausedAndGetVaultState(address(pool));
+    }
+
+    function testVaultAndPoolUnpaused() public {
+        vault.manualSetVaultState(false, true, 3e16, 5e17);
+        vault.manualSetPoolPaused(address(pool), false, 1);
+        VaultState memory vaultState = vault.testEnsureUnpausedAndGetVaultState(address(pool));
+        assertEq(vaultState.isVaultPaused, false);
+        assertEq(vaultState.isQueryDisabled, true);
+        assertEq(vaultState.protocolSwapFeePercentage, 3e16);
+        assertEq(vaultState.protocolYieldFeePercentage, 5e17);
+    }
+}