Governance integration testing (#683)

* Segregate tests and mocks and wire up integration test base

* Switch to a production version of predicting a deployment address

* Add integration test for exchange governor migration

* Add integration test for treassury migration

* Add integration test for migrating the treasury

* Add governance upgrade action to transfer ZRX tokens to new governor

* Add governance upgrade action to transfer wCELO tokens to new governor

* Add governance upgrade action to transfer WYV tokens to new governor

* Turn on verbose logging

contracts/governance/foundry.toml

@@ -2,6 +2,7 @@
 src = 'src'
 out = 'out'
 libs = ['lib', "../utils/contracts/src/"]
+match_path = "test/unit/*.sol"
 fs_permissions = [{ access = "read", path = "./" }]
 remappings = [
     '@openzeppelin/=./lib/openzeppelin-contracts/contracts/',
@@ -12,6 +13,12 @@ solc = '0.8.19'
 optimizer_runs = 20_000
 via_ir = true
 
+[profile.integration]
+match_path = "test/integration/*.sol"
+
+[rpc_endpoints]
+mainnet = "${MAINNET_RPC_URL}"
+
 [profile.smt.model_checker]
 engine = 'chc'
 timeout = 10_000

contracts/governance/package.json

@@ -10,7 +10,8 @@
   "scripts": {
     "test": "forge test",
     "build": "forge build",
-    "build:smt": "FOUNDRY_PROFILE=smt forge build"
+    "build:smt": "FOUNDRY_PROFILE=smt forge build",
+    "test:integration": "source .env && FOUNDRY_PROFILE=integration forge test --fork-url $MAINNET_RPC_URL --fork-block-number 16884148 -vvv"
   },
   "repository": {
     "type": "git",

contracts/governance/test/BaseTest.t.sol

@@ -23,18 +23,13 @@ import "forge-std/Test.sol";
 import "forge-std/console.sol";
 import "@openzeppelin/token/ERC20/ERC20.sol";
 import "@openzeppelin/proxy/ERC1967/ERC1967Proxy.sol";
-import "./ZRXMock.sol";
+import "./mocks/ZRXMock.sol";
 import "../src/ZRXWrappedToken.sol";
 import "../src/ZeroExVotes.sol";
 import "../src/ZeroExTimelock.sol";
 import "../src/ZeroExProtocolGovernor.sol";
 import "../src/ZeroExTreasuryGovernor.sol";
 
-function predict(address deployer, uint256 nonce) pure returns (address) {
-    require(nonce > 0 && nonce < 128);
-    return address(uint160(uint256(keccak256(abi.encodePacked(bytes2(0xd694), deployer, bytes1(uint8(nonce)))))));
-}
-
 contract BaseTest is Test {
     address payable internal account1 = payable(vm.addr(1));
     address payable internal account2 = payable(vm.addr(2));
@@ -54,11 +49,10 @@ contract BaseTest is Test {
         vm.deal(securityCouncil, 1e20);
     }
 
-    function setupGovernance()
+    function setupGovernance(
-        internal
+        IERC20 zrxToken
-        returns (IERC20, ZRXWrappedToken, ZeroExVotes, ZeroExTimelock, ZeroExTimelock, address, address)
+    ) internal returns (ZRXWrappedToken, ZeroExVotes, ZeroExTimelock, ZeroExTimelock, address, address) {
-    {
+        (ZRXWrappedToken token, ZeroExVotes votes) = setupZRXWrappedToken(zrxToken);
-        (IERC20 zrxToken, ZRXWrappedToken token, ZeroExVotes votes) = setupZRXWrappedToken();
 
         vm.startPrank(account1);
         address[] memory proposers = new address[](0);
@@ -86,31 +80,88 @@ contract BaseTest is Test {
         treasuryTimelock.grantRole(treasuryTimelock.CANCELLER_ROLE(), address(treasuryGovernor));
         vm.stopPrank();
 
-        return (
+        return (token, votes, protocolTimelock, treasuryTimelock, address(protocolGovernor), address(treasuryGovernor));
-            zrxToken,
-            token,
-            votes,
-            protocolTimelock,
-            treasuryTimelock,
-            address(protocolGovernor),
-            address(treasuryGovernor)
-        );
     }
 
-    function setupZRXWrappedToken() internal returns (IERC20, ZRXWrappedToken, ZeroExVotes) {
+    function setupZRXWrappedToken(IERC20 zrxToken) internal returns (ZRXWrappedToken, ZeroExVotes) {
         vm.startPrank(account1);
-        bytes memory _bytecode = vm.getCode("./ZRXToken.json");
+        address wTokenPrediction = predictAddress(account1, vm.getNonce(account1) + 2);
-        IERC20 zrxToken;
-        assembly {
-            zrxToken := create(0, add(_bytecode, 0x20), mload(_bytecode))
-        }
-        address wTokenPrediction = predict(account1, vm.getNonce(account1) + 2);
         ZeroExVotes votesImpl = new ZeroExVotes(wTokenPrediction, quadraticThreshold);
         ERC1967Proxy votesProxy = new ERC1967Proxy(address(votesImpl), abi.encodeCall(votesImpl.initialize, ()));
         ZRXWrappedToken wToken = new ZRXWrappedToken(zrxToken, ZeroExVotes(address(votesProxy)));
         vm.stopPrank();
 
         assert(address(wToken) == wTokenPrediction);
-        return (zrxToken, wToken, ZeroExVotes(address(votesProxy)));
+
+        return (wToken, ZeroExVotes(address(votesProxy)));
+    }
+
+    function mockZRXToken() internal returns (IERC20 zrxToken) {
+        vm.startPrank(account1);
+        bytes memory _bytecode = vm.getCode("./ZRXToken.json");
+        assembly {
+            zrxToken := create(0, add(_bytecode, 0x20), mload(_bytecode))
+        }
+        vm.stopPrank();
+    }
+
+    // Sourced from https://github.com/grappafinance/core/blob/master/src/test/utils/Utilities.sol
+    function predictAddress(address _origin, uint256 _nonce) public pure returns (address) {
+        if (_nonce == 0x00) {
+            return
+                address(
+                    uint160(uint256(keccak256(abi.encodePacked(bytes1(0xd6), bytes1(0x94), _origin, bytes1(0x80)))))
+                );
+        }
+        if (_nonce <= 0x7f) {
+            return
+                address(
+                    uint160(uint256(keccak256(abi.encodePacked(bytes1(0xd6), bytes1(0x94), _origin, uint8(_nonce)))))
+                );
+        }
+        if (_nonce <= 0xff) {
+            return
+                address(
+                    uint160(
+                        uint256(
+                            keccak256(
+                                abi.encodePacked(bytes1(0xd7), bytes1(0x94), _origin, bytes1(0x81), uint8(_nonce))
+                            )
+                        )
+                    )
+                );
+        }
+        if (_nonce <= 0xffff) {
+            return
+                address(
+                    uint160(
+                        uint256(
+                            keccak256(
+                                abi.encodePacked(bytes1(0xd8), bytes1(0x94), _origin, bytes1(0x82), uint16(_nonce))
+                            )
+                        )
+                    )
+                );
+        }
+        if (_nonce <= 0xffffff) {
+            return
+                address(
+                    uint160(
+                        uint256(
+                            keccak256(
+                                abi.encodePacked(bytes1(0xd9), bytes1(0x94), _origin, bytes1(0x83), uint24(_nonce))
+                            )
+                        )
+                    )
+                );
+        }
+        return
+            address(
+                uint160(
+                    uint256(
+                        keccak256(abi.encodePacked(bytes1(0xda), bytes1(0x94), _origin, bytes1(0x84), uint32(_nonce)))
+                    )
+                )
+            );
     }
 }

contracts/governance/test/integration/GovernanceE2E.t.sol

@@ -0,0 +1,246 @@
+// SPDX-License-Identifier: Apache-2.0
+/*
+
+  Copyright 2023 ZeroEx Intl.
+
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+
+*/
+
+pragma solidity ^0.8.19;
+
+import "@openzeppelin/token/ERC20/IERC20.sol";
+import "../mocks/IZeroExMock.sol";
+import "../mocks/IZrxTreasuryMock.sol";
+import "../mocks/IStakingMock.sol";
+import "../BaseTest.t.sol";
+import "../../src/ZRXWrappedToken.sol";
+import "../../src/ZeroExVotes.sol";
+import "../../src/ZeroExTimelock.sol";
+import "../../src/ZeroExProtocolGovernor.sol";
+import "../../src/ZeroExTreasuryGovernor.sol";
+
+contract GovernanceE2ETest is BaseTest {
+    uint256 internal mainnetFork;
+    string internal MAINNET_RPC_URL = vm.envString("MAINNET_RPC_URL");
+
+    address internal constant ZRX_TOKEN = 0xE41d2489571d322189246DaFA5ebDe1F4699F498;
+    address internal constant MATIC_TOKEN = 0x7D1AfA7B718fb893dB30A3aBc0Cfc608AaCfeBB0;
+    address internal constant WCELO_TOKEN = 0xE452E6Ea2dDeB012e20dB73bf5d3863A3Ac8d77a;
+    address internal constant WYV_TOKEN = 0x056017c55aE7AE32d12AeF7C679dF83A85ca75Ff;
+
+    address internal constant EXCHANGE_PROXY = 0xDef1C0ded9bec7F1a1670819833240f027b25EfF;
+    address internal constant EXCHANGE_GOVERNOR = 0x618F9C67CE7Bf1a50afa1E7e0238422601b0ff6e;
+    address internal constant TREASURY = 0x0bB1810061C2f5b2088054eE184E6C79e1591101;
+    address internal constant STAKING = 0xa26e80e7Dea86279c6d778D702Cc413E6CFfA777;
+    address internal staker = 0x5265Bde27F57E738bE6c1F6AB3544e82cdc92a8f;
+    bytes32 internal stakerPool = 0x0000000000000000000000000000000000000000000000000000000000000032;
+    bytes32[] internal staker_operated_poolIds = [stakerPool];
+
+    // voting power 1500000e18
+    address internal voter1 = 0x292c6DAE7417B3D31d8B6e1d2EeA0258d14C4C4b;
+    bytes32 internal voter1Pool = 0x0000000000000000000000000000000000000000000000000000000000000030;
+    bytes32[] internal voter1_operated_poolIds = [voter1Pool];
+
+    // voting power 1500000.5e18
+    address internal voter2 = 0x4990cE223209FCEc4ec4c1ff6E0E81eebD8Cca08;
+    bytes32 internal voter2Pool = 0x0000000000000000000000000000000000000000000000000000000000000031;
+    bytes32[] internal voter2_operated_poolIds = [voter2Pool];
+
+    // voting power 1500000e18
+    address internal voter3 = 0x5265Bde27F57E738bE6c1F6AB3544e82cdc92a8f;
+    bytes32 internal voter3Pool = 0x0000000000000000000000000000000000000000000000000000000000000032;
+    bytes32[] internal voter3_operated_poolIds = [voter3Pool];
+
+    // voting power 1500000e18
+    address internal voter4 = 0xcA9F5049c1Ea8FC78574f94B7Cf5bE5fEE354C31;
+    bytes32 internal voter4Pool = 0x0000000000000000000000000000000000000000000000000000000000000034;
+    bytes32[] internal voter4_operated_poolIds = [voter4Pool];
+
+    // voting power 1500000e18
+    address internal voter5 = 0xDBB5664a9DBCB98F6365804880e5b277B3155422;
+    bytes32 internal voter5Pool = 0x0000000000000000000000000000000000000000000000000000000000000035;
+    bytes32[] internal voter5_operated_poolIds = [voter5Pool];
+
+    // voting power 2291490.952353335e18
+    address internal voter6 = 0x9a4Eb1101C0c053505Bd71d2fFa27Ed902DEaD85;
+    bytes32 internal voter6Pool = 0x0000000000000000000000000000000000000000000000000000000000000029;
+    bytes32[] internal voter6_operated_poolIds = [voter6Pool];
+
+    // voting power 4575984.325e18
+    address internal voter7 = 0x9564177EC8052C92752a488a71769F710aA0A41D;
+    bytes32 internal voter7Pool = 0x0000000000000000000000000000000000000000000000000000000000000025;
+    bytes32[] internal voter7_operated_poolIds = [voter7Pool];
+
+    IERC20 internal token;
+    IERC20 internal maticToken;
+    IERC20 internal wceloToken;
+    IERC20 internal wyvToken;
+
+    IZeroExMock internal exchange;
+    IZrxTreasuryMock internal treasury;
+    IStakingMock internal staking;
+
+    ZRXWrappedToken internal wToken;
+    ZeroExVotes internal votes;
+    ZeroExTimelock internal protocolTimelock;
+    ZeroExTimelock internal treasuryTimelock;
+    ZeroExProtocolGovernor internal protocolGovernor;
+    ZeroExTreasuryGovernor internal treasuryGovernor;
+
+    function setUp() public {
+        mainnetFork = vm.createFork(MAINNET_RPC_URL);
+        vm.selectFork(mainnetFork);
+
+        token = IERC20(ZRX_TOKEN);
+        maticToken = IERC20(MATIC_TOKEN);
+        wceloToken = IERC20(WCELO_TOKEN);
+        wyvToken = IERC20(WYV_TOKEN);
+
+        exchange = IZeroExMock(payable(EXCHANGE_PROXY));
+        treasury = IZrxTreasuryMock(TREASURY);
+        staking = IStakingMock(STAKING);
+
+        address protocolGovernorAddress;
+        address treasuryGovernorAddress;
+        (
+            wToken,
+            votes,
+            protocolTimelock,
+            treasuryTimelock,
+            protocolGovernorAddress,
+            treasuryGovernorAddress
+        ) = setupGovernance(token);
+
+        protocolGovernor = ZeroExProtocolGovernor(payable(protocolGovernorAddress));
+        treasuryGovernor = ZeroExTreasuryGovernor(payable(treasuryGovernorAddress));
+    }
+
+    function testProtocolGovernanceMigration() public {
+        // initially the zrx exchange is owned by the legacy exchange governor
+        assertEq(exchange.owner(), EXCHANGE_GOVERNOR);
+
+        // transfer ownership to new protocol governor
+        vm.prank(EXCHANGE_GOVERNOR);
+        exchange.transferOwnership(address(protocolGovernor));
+        assertEq(exchange.owner(), address(protocolGovernor));
+    }
+
+    function testTreasuryGovernanceMigration() public {
+        // Create a proposal to migrate to new governor
+
+        uint256 currentEpoch = staking.currentEpoch();
+        uint256 executionEpoch = currentEpoch + 2;
+
+        vm.startPrank(staker);
+
+        IZrxTreasuryMock.ProposedAction[] memory actions = new IZrxTreasuryMock.ProposedAction[](4);
+
+        // Transfer MATIC
+        uint256 maticBalance = maticToken.balanceOf(address(treasury));
+        actions[0] = IZrxTreasuryMock.ProposedAction({
+            target: MATIC_TOKEN,
+            data: abi.encodeCall(maticToken.transfer, (address(treasuryGovernor), maticBalance)),
+            value: 0
+        });
+
+        // Transfer ZRX
+        uint256 zrxBalance = token.balanceOf(address(treasury));
+        actions[1] = IZrxTreasuryMock.ProposedAction({
+            target: ZRX_TOKEN,
+            data: abi.encodeCall(token.transfer, (address(treasuryGovernor), zrxBalance)),
+            value: 0
+        });
+
+        // Transfer wCELO
+        uint256 wceloBalance = wceloToken.balanceOf(address(treasury));
+        actions[2] = IZrxTreasuryMock.ProposedAction({
+            target: WCELO_TOKEN,
+            data: abi.encodeCall(wceloToken.transfer, (address(treasuryGovernor), wceloBalance)),
+            value: 0
+        });
+
+        // Transfer WYV
+        uint256 wyvBalance = wyvToken.balanceOf(address(treasury));
+        actions[3] = IZrxTreasuryMock.ProposedAction({
+            target: WYV_TOKEN,
+            data: abi.encodeCall(wyvToken.transfer, (address(treasuryGovernor), wyvBalance)),
+            value: 0
+        });
+
+        uint256 proposalId = treasury.propose(
+            actions,
+            executionEpoch,
+            "Z-5 Migrate to new treasury governor",
+            staker_operated_poolIds
+        );
+
+        // Once a proposal is created, it becomes open for voting at the epoch after next (currentEpoch + 2)
+        // and is open for the voting period (currently set to 3 days).
+        uint256 epochDurationInSeconds = staking.epochDurationInSeconds(); // Currently set to 604800 seconds = 7 days
+        uint256 currentEpochEndTime = staking.currentEpochStartTimeInSeconds() + epochDurationInSeconds;
+
+        vm.warp(currentEpochEndTime + 1);
+        staking.endEpoch();
+        vm.warp(block.timestamp + epochDurationInSeconds + 1);
+        staking.endEpoch();
+
+        vm.stopPrank();
+        // quorum is 10,000,000e18 so reach that via the following votes
+        vm.prank(voter1);
+        treasury.castVote(proposalId, true, voter1_operated_poolIds);
+        vm.stopPrank();
+
+        vm.prank(voter2);
+        treasury.castVote(proposalId, true, voter2_operated_poolIds);
+        vm.stopPrank();
+
+        vm.prank(voter3);
+        treasury.castVote(proposalId, true, voter3_operated_poolIds);
+        vm.stopPrank();
+
+        vm.prank(voter4);
+        treasury.castVote(proposalId, true, voter4_operated_poolIds);
+        vm.stopPrank();
+
+        vm.prank(voter5);
+        treasury.castVote(proposalId, true, voter5_operated_poolIds);
+        vm.stopPrank();
+
+        vm.prank(voter6);
+        treasury.castVote(proposalId, true, voter6_operated_poolIds);
+        vm.stopPrank();
+
+        vm.prank(voter7);
+        treasury.castVote(proposalId, true, voter7_operated_poolIds);
+        vm.stopPrank();
+
+        vm.warp(block.timestamp + 3 days + 1);
+
+        // Execute proposal
+        treasury.execute(proposalId, actions);
+
+        // Assert value of treasury has correctly transferred
+        uint256 maticBalanceNewTreasury = maticToken.balanceOf(address(treasuryGovernor));
+        assertEq(maticBalanceNewTreasury, maticBalance);
+
+        uint256 zrxBalanceNewTreasury = token.balanceOf(address(treasuryGovernor));
+        assertEq(zrxBalanceNewTreasury, zrxBalance);
+
+        uint256 wceloBalanceNewTreasury = wceloToken.balanceOf(address(treasuryGovernor));
+        assertEq(wceloBalanceNewTreasury, wceloBalance);
+
+        uint256 wyvBalanceNewTreasury = wyvToken.balanceOf(address(treasuryGovernor));
+        assertEq(wyvBalanceNewTreasury, wyvBalance);
+    }
+}

contracts/governance/test/mocks/IOwnableFeature.sol

@@ -0,0 +1,36 @@
+// SPDX-License-Identifier: Apache-2.0
+/*
+  Copyright 2023 ZeroEx Intl.
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+*/
+
+pragma solidity ^0.8.19;
+
+import "@0x/contracts-utils/contracts/src/v08/interfaces/IOwnableV08.sol";
+
+/// @dev Owner management and migration features.
+interface IOwnableFeature is IOwnableV08 {
+    /// @dev Emitted when `migrate()` is called.
+    /// @param caller The caller of `migrate()`.
+    /// @param migrator The migration contract.
+    /// @param newOwner The address of the new owner.
+    event Migrated(address caller, address migrator, address newOwner);
+
+    /// @dev Execute a migration function in the context of the ZeroEx contract.
+    ///      The result of the function being called should be the magic bytes
+    ///      0x2c64c5ef (`keccack('MIGRATE_SUCCESS')`). Only callable by the owner.
+    ///      The owner will be temporarily set to `address(this)` inside the call.
+    ///      Before returning, the owner will be set to `newOwner`.
+    /// @param target The migrator contract address.
+    /// @param newOwner The address of the new owner.
+    /// @param data The call data.
+    function migrate(address target, bytes calldata data, address newOwner) external;
+}

contracts/governance/test/mocks/ISimpleFunctionRegistryFeature.sol

@@ -0,0 +1,48 @@
+// SPDX-License-Identifier: Apache-2.0
+/*
+  Copyright 2023 ZeroEx Intl.
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+*/
+
+pragma solidity ^0.8.19;
+pragma experimental ABIEncoderV2;
+
+/// @dev Basic registry management features.
+interface ISimpleFunctionRegistryFeature {
+    /// @dev A function implementation was updated via `extend()` or `rollback()`.
+    /// @param selector The function selector.
+    /// @param oldImpl The implementation contract address being replaced.
+    /// @param newImpl The replacement implementation contract address.
+    event ProxyFunctionUpdated(bytes4 indexed selector, address oldImpl, address newImpl);
+
+    /// @dev Roll back to a prior implementation of a function.
+    /// @param selector The function selector.
+    /// @param targetImpl The address of an older implementation of the function.
+    function rollback(bytes4 selector, address targetImpl) external;
+
+    /// @dev Register or replace a function.
+    /// @param selector The function selector.
+    /// @param impl The implementation contract for the function.
+    function extend(bytes4 selector, address impl) external;
+
+    /// @dev Retrieve the length of the rollback history for a function.
+    /// @param selector The function selector.
+    /// @return rollbackLength The number of items in the rollback history for
+    ///         the function.
+    function getRollbackLength(bytes4 selector) external view returns (uint256 rollbackLength);
+
+    /// @dev Retrieve an entry in the rollback history for a function.
+    /// @param selector The function selector.
+    /// @param idx The index in the rollback history.
+    /// @return impl An implementation address for the function at
+    ///         index `idx`.
+    function getRollbackEntryAtIndex(bytes4 selector, uint256 idx) external view returns (address impl);
+}

contracts/governance/test/mocks/IStakingMock.sol

@@ -0,0 +1,106 @@
+// SPDX-License-Identifier: Apache-2.0
+/*
+
+  Copyright 2021 ZeroEx Intl.
+
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+
+*/
+
+pragma solidity ^0.8.19;
+
+interface IStakingMock {
+    /// @dev Statuses that stake can exist in.
+    ///      Any stake can be (re)delegated effective at the next epoch
+    ///      Undelegated stake can be withdrawn if it is available in both the current and next epoch
+    enum StakeStatus {
+        UNDELEGATED,
+        DELEGATED
+    }
+
+    /// @dev Encapsulates a balance for the current and next epochs.
+    /// Note that these balances may be stale if the current epoch
+    /// is greater than `currentEpoch`.
+    /// @param currentEpoch The current epoch
+    /// @param currentEpochBalance Balance in the current epoch.
+    /// @param nextEpochBalance Balance in `currentEpoch+1`.
+    struct StoredBalance {
+        uint64 currentEpoch;
+        uint96 currentEpochBalance;
+        uint96 nextEpochBalance;
+    }
+
+    /// @dev Holds the metadata for a staking pool.
+    /// @param operator Operator of the pool.
+    /// @param operatorShare Fraction of the total balance owned by the operator, in ppm.
+    struct Pool {
+        address operator;
+        uint32 operatorShare;
+    }
+
+    /// @dev Create a new staking pool. The sender will be the operator of this pool.
+    /// Note that an operator must be payable.
+    /// @param operatorShare Portion of rewards owned by the operator, in ppm.
+    /// @param addOperatorAsMaker Adds operator to the created pool as a maker for convenience iff true.
+    /// @return poolId The unique pool id generated for this pool.
+    function createStakingPool(uint32 operatorShare, bool addOperatorAsMaker) external returns (bytes32 poolId);
+
+    /// @dev Returns the current staking epoch number.
+    /// @return epoch The current epoch.
+    function currentEpoch() external view returns (uint256 epoch);
+
+    /// @dev Returns the time (in seconds) at which the current staking epoch started.
+    /// @return startTime The start time of the current epoch, in seconds.
+    function currentEpochStartTimeInSeconds() external view returns (uint256 startTime);
+
+    /// @dev Returns the duration of an epoch in seconds. This value can be updated.
+    /// @return duration The duration of an epoch, in seconds.
+    function epochDurationInSeconds() external view returns (uint256 duration);
+
+    /// @dev Returns a staking pool
+    /// @param poolId Unique id of pool.
+    function getStakingPool(bytes32 poolId) external view returns (Pool memory);
+
+    /// @dev Gets global stake for a given status.
+    /// @param stakeStatus UNDELEGATED or DELEGATED
+    /// @return balance Global stake for given status.
+    function getGlobalStakeByStatus(StakeStatus stakeStatus) external view returns (StoredBalance memory balance);
+
+    /// @dev Gets an owner's stake balances by status.
+    /// @param staker Owner of stake.
+    /// @param stakeStatus UNDELEGATED or DELEGATED
+    /// @return balance Owner's stake balances for given status.
+    function getOwnerStakeByStatus(
+        address staker,
+        StakeStatus stakeStatus
+    ) external view returns (StoredBalance memory balance);
+
+    /// @dev Returns the total stake delegated to a specific staking pool,
+    ///      across all members.
+    /// @param poolId Unique Id of pool.
+    /// @return balance Total stake delegated to pool.
+    function getTotalStakeDelegatedToPool(bytes32 poolId) external view returns (StoredBalance memory balance);
+
+    /// @dev Returns the stake delegated to a specific staking pool, by a given staker.
+    /// @param staker of stake.
+    /// @param poolId Unique Id of pool.
+    /// @return balance Stake delegated to pool by staker.
+    function getStakeDelegatedToPoolByOwner(
+        address staker,
+        bytes32 poolId
+    ) external view returns (StoredBalance memory balance);
+
+    function endEpoch() external returns (uint256);
+
+    function finalizePool(bytes32 poolId) external;
+}

contracts/governance/test/mocks/IZeroExMock.sol

@@ -0,0 +1,24 @@
+// SPDX-License-Identifier: Apache-2.0
+/*
+  Copyright 2023 ZeroEx Intl.
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+*/
+
+pragma solidity ^0.8.19;
+
+import "./IOwnableFeature.sol";
+import "./ISimpleFunctionRegistryFeature.sol";
+
+/// @dev Minimal viable Exchange Proxy interface for governance use.
+interface IZeroExMock is IOwnableFeature, ISimpleFunctionRegistryFeature {
+    /// @dev Fallback for just receiving ether.
+    receive() external payable;
+}

contracts/governance/test/mocks/IZrxTreasuryMock.sol

@@ -0,0 +1,159 @@
+// SPDX-License-Identifier: Apache-2.0
+/*
+
+  Copyright 2021 ZeroEx Intl.
+
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+
+*/
+
+pragma solidity ^0.8.19;
+
+import "./IStakingMock.sol";
+
+/// @dev Minimal viable Treasury interface for governance use.
+interface IZrxTreasuryMock {
+    struct TreasuryParameters {
+        uint256 votingPeriod;
+        uint256 proposalThreshold;
+        uint256 quorumThreshold;
+        bytes32 defaultPoolId;
+    }
+
+    struct ProposedAction {
+        address target;
+        bytes data;
+        uint256 value;
+    }
+
+    struct Proposal {
+        bytes32 actionsHash;
+        uint256 executionEpoch;
+        uint256 voteEpoch;
+        uint256 votesFor;
+        uint256 votesAgainst;
+        bool executed;
+    }
+
+    event ProposalCreated(
+        address proposer,
+        bytes32[] operatedPoolIds,
+        uint256 proposalId,
+        ProposedAction[] actions,
+        uint256 executionEpoch,
+        string description
+    );
+
+    event VoteCast(address voter, bytes32[] operatedPoolIds, uint256 proposalId, bool support, uint256 votingPower);
+
+    event ProposalExecuted(uint256 proposalId);
+
+    function stakingProxy() external view returns (IStakingMock);
+
+    function defaultPoolId() external view returns (bytes32);
+
+    function votingPeriod() external view returns (uint256);
+
+    function proposalThreshold() external view returns (uint256);
+
+    function quorumThreshold() external view returns (uint256);
+
+    /// @dev Updates the proposal and quorum thresholds to the given
+    ///      values. Note that this function is only callable by the
+    ///      treasury contract itself, so the threshold can only be
+    ///      updated via a successful treasury proposal.
+    /// @param newProposalThreshold The new value for the proposal threshold.
+    /// @param newQuorumThreshold The new value for the quorum threshold.
+    function updateThresholds(uint256 newProposalThreshold, uint256 newQuorumThreshold) external;
+
+    /// @dev Creates a proposal to send ZRX from this treasury on the
+    ///      the given actions. Must have at least `proposalThreshold`
+    ///      of voting power to call this function. See `getVotingPower`
+    ///      for how voting power is computed. If a proposal is successfully
+    ///      created, voting starts at the epoch after next (currentEpoch + 2).
+    ///      If the vote passes, the proposal is executable during the
+    ///      `executionEpoch`. See `hasProposalPassed` for the passing criteria.
+    /// @param actions The proposed ZRX actions. An action specifies a
+    ///        contract call.
+    /// @param executionEpoch The epoch during which the proposal is to
+    ///        be executed if it passes. Must be at least two epochs
+    ///        from the current epoch.
+    /// @param description A text description for the proposal.
+    /// @param operatedPoolIds The pools operated by `msg.sender`. The
+    ///        ZRX currently delegated to those pools will be accounted
+    ///        for in the voting power.
+    /// @return proposalId The ID of the newly created proposal.
+    function propose(
+        ProposedAction[] calldata actions,
+        uint256 executionEpoch,
+        string calldata description,
+        bytes32[] calldata operatedPoolIds
+    ) external returns (uint256 proposalId);
+
+    /// @dev Casts a vote for the given proposal. Only callable
+    ///      during the voting period for that proposal.
+    ///      One address can only vote once.
+    ///      See `getVotingPower` for how voting power is computed.
+    /// @param proposalId The ID of the proposal to vote on.
+    /// @param support Whether to support the proposal or not.
+    /// @param operatedPoolIds The pools operated by `msg.sender`. The
+    ///        ZRX currently delegated to those pools will be accounted
+    ///        for in the voting power.
+    function castVote(uint256 proposalId, bool support, bytes32[] calldata operatedPoolIds) external;
+
+    /// @dev Casts a vote for the given proposal, by signature.
+    ///      Only callable during the voting period for that proposal.
+    ///      One address/voter can only vote once.
+    ///      See `getVotingPower` for how voting power is computed.
+    /// @param proposalId The ID of the proposal to vote on.
+    /// @param support Whether to support the proposal or not.
+    /// @param operatedPoolIds The pools operated by the signer. The
+    ///        ZRX currently delegated to those pools will be accounted
+    ///        for in the voting power.
+    /// @param v the v field of the signature
+    /// @param r the r field of the signature
+    /// @param s the s field of the signature
+    function castVoteBySignature(
+        uint256 proposalId,
+        bool support,
+        bytes32[] memory operatedPoolIds,
+        uint8 v,
+        bytes32 r,
+        bytes32 s
+    ) external;
+
+    /// @dev Executes a proposal that has passed and is
+    ///      currently executable.
+    /// @param proposalId The ID of the proposal to execute.
+    /// @param actions Actions associated with the proposal to execute.
+    function execute(uint256 proposalId, ProposedAction[] memory actions) external payable;
+
+    /// @dev Returns the total number of proposals.
+    /// @return count The number of proposals.
+    function proposalCount() external view returns (uint256 count);
+
+    /// @dev Computes the current voting power of the given account.
+    ///      Voting power is equal to:
+    ///      (ZRX delegated to the default pool) +
+    ///      0.5 * (ZRX delegated to other pools) +
+    ///      0.5 * (ZRX delegated to pools operated by account)
+    /// @param account The address of the account.
+    /// @param operatedPoolIds The pools operated by `account`. The
+    ///        ZRX currently delegated to those pools will be accounted
+    ///        for in the voting power.
+    /// @return votingPower The current voting power of the given account.
+    function getVotingPower(
+        address account,
+        bytes32[] calldata operatedPoolIds
+    ) external view returns (uint256 votingPower);
+}

contracts/governance/test/mocks/ZeroExVotesMalicious.sol

@@ -18,7 +18,7 @@
 */
 pragma solidity ^0.8.19;
 
-import "../src/ZeroExVotes.sol";
+import "../../src/ZeroExVotes.sol";
 
 contract ZeroExVotesMalicious is ZeroExVotes {
     constructor(address _token, uint256 _quadraticThreshold) ZeroExVotes(_token, _quadraticThreshold) {}

contracts/governance/test/mocks/ZeroExVotesMigration.sol

@@ -18,7 +18,7 @@
 */
 pragma solidity ^0.8.19;
 
-import {ZeroExVotes} from "../src/ZeroExVotes.sol";
+import {ZeroExVotes} from "../../src/ZeroExVotes.sol";
 import {SafeCast} from "@openzeppelin/utils/math/SafeCast.sol";
 import {Math} from "@openzeppelin/utils/math/Math.sol";
 import {CubeRoot} from "./CubeRoot.sol";

contracts/governance/test/unit/ZRXWrappedTokenTest.t.sol

@@ -18,8 +18,8 @@
 */
 pragma solidity ^0.8.19;
 
-import "./BaseTest.t.sol";
+import "../BaseTest.t.sol";
-import "../src/ZRXWrappedToken.sol";
+import "../../src/ZRXWrappedToken.sol";
 import "@openzeppelin/token/ERC20/ERC20.sol";
 
 contract ZRXWrappedTokenTest is BaseTest {
@@ -28,7 +28,8 @@ contract ZRXWrappedTokenTest is BaseTest {
     ZeroExVotes private votes;
 
     function setUp() public {
-        (token, wToken, votes, , , , ) = setupGovernance();
+        token = mockZRXToken();
+        (wToken, votes, , , , ) = setupGovernance(token);
         vm.startPrank(account1);
         token.transfer(account2, 100e18);
         token.transfer(account3, 200e18);

contracts/governance/test/unit/ZeroExGovernorBaseTest.t.sol

@@ -18,11 +18,11 @@
 */
 pragma solidity ^0.8.19;
 
-import "./BaseTest.t.sol";
+import "../BaseTest.t.sol";
-import "../src/IZeroExGovernor.sol";
+import "../../src/IZeroExGovernor.sol";
-import "../src/ZeroExTimelock.sol";
+import "../../src/ZeroExTimelock.sol";
-import "../src/ZeroExProtocolGovernor.sol";
+import "../../src/ZeroExProtocolGovernor.sol";
-import "../src/ZRXWrappedToken.sol";
+import "../../src/ZRXWrappedToken.sol";
 import "@openzeppelin/token/ERC20/ERC20.sol";
 import "@openzeppelin/mocks/CallReceiverMock.sol";
 

contracts/governance/test/unit/ZeroExProtocolGovernor.t.sol

@@ -19,8 +19,8 @@
 pragma solidity ^0.8.19;
 
 import "./ZeroExGovernorBaseTest.t.sol";
-import "./ZeroExMock.sol";
+import "../mocks/ZeroExMock.sol";
-import "../src/ZeroExProtocolGovernor.sol";
+import "../../src/ZeroExProtocolGovernor.sol";
 
 contract ZeroExProtocolGovernorTest is ZeroExGovernorBaseTest {
     ZeroExProtocolGovernor internal protocolGovernor;
@@ -35,7 +35,9 @@ contract ZeroExProtocolGovernorTest is ZeroExGovernorBaseTest {
         quorum = 10000000e18;
 
         address governorAddress;
-        (token, wToken, votes, timelock, , governorAddress, ) = setupGovernance();
+
+        token = mockZRXToken();
+        (wToken, votes, timelock, , governorAddress, ) = setupGovernance(token);
         governor = IZeroExGovernor(governorAddress);
         protocolGovernor = ZeroExProtocolGovernor(payable(governorAddress));
         zeroExMock = new ZeroExMock();

contracts/governance/test/unit/ZeroExTreasuryGovernor.t.sol

@@ -26,7 +26,8 @@ contract ZeroExTreasuryGovernorTest is ZeroExGovernorBaseTest {
         proposalThreshold = 250000e18;
 
         address governorAddress;
-        (token, wToken, votes, , timelock, , governorAddress) = setupGovernance();
+        token = mockZRXToken();
+        (wToken, votes, , timelock, , governorAddress) = setupGovernance(token);
         governor = IZeroExGovernor(governorAddress);
 
         initialiseAccounts();

contracts/governance/test/unit/ZeroExVotesTest.t.sol

@@ -19,18 +19,20 @@
 pragma solidity ^0.8.19;
 
 import "@openzeppelin/token/ERC20/ERC20.sol";
-import "./BaseTest.t.sol";
+import "../BaseTest.t.sol";
-import "./ZeroExVotesMalicious.sol";
+import "../mocks/ZeroExVotesMalicious.sol";
-import "./ZeroExVotesMigration.sol";
+import "../mocks/ZeroExVotesMigration.sol";
-import "../src/ZRXWrappedToken.sol";
+import "../../src/ZRXWrappedToken.sol";
+import "../../src/ZeroExVotes.sol";
 
 contract ZeroExVotesTest is BaseTest {
-    IERC20 private token;
+    IERC20 internal token;
-    ZRXWrappedToken private wToken;
+    ZRXWrappedToken internal wToken;
-    ZeroExVotes private votes;
+    ZeroExVotes internal votes;
 
     function setUp() public {
-        (token, wToken, votes) = setupZRXWrappedToken();
+        token = mockZRXToken();
+        (wToken, votes) = setupZRXWrappedToken(token);
         vm.startPrank(account1);
         token.transfer(account2, 1700000e18);
         token.transfer(account3, 1600000e18);