@0x/contracts-staking: Updating tests and making the contracts testable.

contracts/staking/contracts/src/fees/MixinExchangeFees.sol

@@ -192,24 +192,23 @@ contract MixinExchangeFees is
         private
         view
     {
-        if (protocolFeePaid != 0) {
+        if (protocolFeePaid == 0) {
-            return;
+            LibRichErrors.rrevert(
+                LibStakingRichErrors.InvalidProtocolFeePaymentError(
+                    LibStakingRichErrors.ProtocolFeePaymentErrorCodes.ZeroProtocolFeePaid,
+                    protocolFeePaid,
+                    msg.value
+                )
+            );
         }
-        if (msg.value == protocolFeePaid || msg.value == 0) {
+        if (msg.value != protocolFeePaid && msg.value != 0) {
-            return;
+            LibRichErrors.rrevert(
+                LibStakingRichErrors.InvalidProtocolFeePaymentError(
+                    LibStakingRichErrors.ProtocolFeePaymentErrorCodes.MismatchedFeeAndPayment,
+                    protocolFeePaid,
+                    msg.value
+                )
+            );
         }
-        LibRichErrors.rrevert(
-            LibStakingRichErrors.InvalidProtocolFeePaymentError(
-                protocolFeePaid == 0 ?
-                    LibStakingRichErrors
-                        .ProtocolFeePaymentErrorCodes
-                        .ZeroProtocolFeePaid :
-                    LibStakingRichErrors
-                        .ProtocolFeePaymentErrorCodes
-                        .MismatchedFeeAndPayment,
-                protocolFeePaid,
-                msg.value
-            )
-        );
     }
 }

contracts/staking/contracts/src/sys/MixinAbstract.sol

@@ -26,20 +26,6 @@ import "../interfaces/IStructs.sol";
 ///      cyclical dependencies.
 contract MixinAbstract {
 
-    /// @dev Computes the reward owed to a pool during finalization.
-    ///      Does nothing if the pool is already finalized.
-    /// @param poolId The pool's ID.
-    /// @return totalReward The total reward owed to a pool.
-    /// @return membersStake The total stake for all non-operator members in
-    ///         this pool.
-    function _getUnfinalizedPoolRewards(bytes32 poolId)
-        internal
-        view
-        returns (
-            uint256 totalReward,
-            uint256 membersStake
-        );
-
     /// @dev Instantly finalizes a single pool that was active in the previous
     ///      epoch, crediting it rewards and sending those rewards to the reward
     ///      and eth vault. This can be called by internal functions that need
@@ -57,4 +43,18 @@ contract MixinAbstract {
             uint256 membersReward,
             uint256 membersStake
         );
+
+    /// @dev Computes the reward owed to a pool during finalization.
+    ///      Does nothing if the pool is already finalized.
+    /// @param poolId The pool's ID.
+    /// @return totalReward The total reward owed to a pool.
+    /// @return membersStake The total stake for all non-operator members in
+    ///         this pool.
+    function _getUnfinalizedPoolRewards(bytes32 poolId)
+        internal
+        view
+        returns (
+            uint256 totalReward,
+            uint256 membersStake
+        );
 }

contracts/staking/contracts/src/sys/MixinFinalizer.sol

@@ -248,7 +248,7 @@ contract MixinFinalizer is
         IEtherToken weth = IEtherToken(_getWETHAddress());
         uint256 ethBalance = address(this).balance;
         if (ethBalance != 0) {
-            weth.deposit.value((address(this).balance));
+            weth.deposit.value((address(this).balance))();
         }
         balance = weth.balanceOf(address(this));
         return balance;

contracts/staking/contracts/src/sys/MixinParams.sol

@@ -235,7 +235,7 @@ contract MixinParams is
         address[2] memory oldSpenders,
         address[2] memory newSpenders
     )
-        private
+        internal
     {
         IEtherToken weth = IEtherToken(_getWETHAddress());
         // Grant new allowances.

contracts/staking/contracts/src/vaults/EthVault.sol

@@ -21,7 +21,6 @@ pragma solidity ^0.5.9;
 import "@0x/contracts-erc20/contracts/src/interfaces/IEtherToken.sol";
 import "@0x/contracts-utils/contracts/src/LibSafeMath.sol";
 import "../interfaces/IEthVault.sol";
-import "../immutable/MixinDeploymentConstants.sol";
 import "./MixinVaultCore.sol";
 
 
@@ -29,15 +28,21 @@ import "./MixinVaultCore.sol";
 contract EthVault is
     IEthVault,
     IVaultCore,
-    MixinDeploymentConstants,
     Ownable,
     MixinVaultCore
 {
     using LibSafeMath for uint256;
 
+    // Address of the WETH contract.
+    IEtherToken public weth;
     // mapping from Owner to WETH balance
     mapping (address => uint256) internal _balances;
 
+    /// @param wethAddress Address of the WETH contract.
+    constructor(address wethAddress) public {
+        weth = IEtherToken(wethAddress);
+    }
+
     /// @dev Deposit an `amount` of WETH for `owner` into the vault.
     ///      The staking contract should have granted the vault an allowance
     ///      because it will pull the WETH via `transferFrom()`.
@@ -49,7 +54,7 @@ contract EthVault is
         onlyStakingProxy
     {
         // Transfer WETH from the staking contract into this contract.
-        IEtherToken(_getWETHAddress()).transferFrom(msg.sender, address(this), amount);
+        weth.transferFrom(msg.sender, address(this), amount);
         // Credit the owner.
         _balances[owner] = _balances[owner].safeAdd(amount);
         emit EthDepositedIntoVault(msg.sender, owner, amount);
@@ -97,7 +102,7 @@ contract EthVault is
         _balances[owner] = _balances[owner].safeSub(amount);
 
         // withdraw WETH to owner
-        IEtherToken(_getWETHAddress()).transfer(msg.sender, amount);
+        weth.transfer(msg.sender, amount);
 
         // notify
         emit EthWithdrawnFromVault(msg.sender, owner, amount);

contracts/staking/contracts/src/vaults/StakingPoolRewardVault.sol

@@ -26,22 +26,27 @@ import "../libs/LibStakingRichErrors.sol";
 import "../libs/LibSafeDowncast.sol";
 import "./MixinVaultCore.sol";
 import "../interfaces/IStakingPoolRewardVault.sol";
-import "../immutable/MixinDeploymentConstants.sol";
 
 
 /// @dev This vault manages staking pool rewards.
 contract StakingPoolRewardVault is
     IStakingPoolRewardVault,
     IVaultCore,
-    MixinDeploymentConstants,
     Ownable,
     MixinVaultCore
 {
     using LibSafeMath for uint256;
 
+    // Address of the WETH contract.
+    IEtherToken public weth;
     // mapping from poolId to Pool metadata
     mapping (bytes32 => uint256) internal _balanceByPoolId;
 
+    /// @param wethAddress Address of the WETH contract.
+    constructor(address wethAddress) public {
+        weth = IEtherToken(wethAddress);
+    }
+
     /// @dev Deposit an amount of WETH for `poolId` into the vault.
     ///      The staking contract should have granted the vault an allowance
     ///      because it will pull the WETH via `transferFrom()`.
@@ -53,7 +58,7 @@ contract StakingPoolRewardVault is
         onlyStakingProxy
     {
         // Transfer WETH from the staking contract into this contract.
-        IEtherToken(_getWETHAddress()).transferFrom(msg.sender, address(this), amount);
+        weth.transferFrom(msg.sender, address(this), amount);
         // Credit the pool.
         _balanceByPoolId[poolId] = _balanceByPoolId[poolId].safeAdd(amount);
         emit EthDepositedIntoVault(msg.sender, poolId, amount);
@@ -73,7 +78,7 @@ contract StakingPoolRewardVault is
         onlyStakingProxy
     {
         _balanceByPoolId[poolId] = _balanceByPoolId[poolId].safeSub(amount);
-        IEtherToken(_getWETHAddress()).transfer(to, amount);
+        weth.transfer(to, amount);
         emit PoolRewardTransferred(
             poolId,
             to,

contracts/staking/contracts/test/TestCumulativeRewardTracking.sol

@@ -21,6 +21,7 @@ pragma experimental ABIEncoderV2;
 import "./TestStaking.sol";
 
 
+// solhint-disable no-empty-blocks
 contract TestCumulativeRewardTracking is
     TestStaking
 {
@@ -39,7 +40,8 @@ contract TestCumulativeRewardTracking is
         uint256 epoch
     );
 
-    // solhint-disable-next-line no-empty-blocks
+    constructor(address wethAddress) public TestStaking(wethAddress) {}
+
     function init(address, address, address payable, address) public {}
 
     function _forceSetCumulativeReward(

contracts/staking/contracts/test/TestMixinParams.sol

@@ -0,0 +1,54 @@
+/*
+
+  Copyright 2019 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.5.9;
+pragma experimental ABIEncoderV2;
+
+import "../src/interfaces/IEthVault.sol";
+import "../src/interfaces/IStakingPoolRewardVault.sol";
+import "../src/sys/MixinParams.sol";
+
+
+// solhint-disable no-empty-blocks
+contract TestMixinParams is
+    MixinParams
+{
+
+    event WETHApprove(address spender, uint256 amount);
+
+    /// @dev Sets the eth and reward vault addresses.
+    function setVaultAddresses(
+        address ethVaultAddress,
+        address rewardVaultAddress
+    )
+        external
+    {
+        ethVault = IEthVault(ethVaultAddress);
+        rewardVault = IStakingPoolRewardVault(rewardVaultAddress);
+    }
+
+    /// @dev WETH `approve()` function that just logs events.
+    function approve(address spender, uint256 amount) external returns (bool) {
+        emit WETHApprove(spender, amount);
+    }
+
+    /// @dev Overridden return this contract's address.
+    function _getWETHAddress() internal view returns (address) {
+        return address(this);
+    }
+}

contracts/staking/contracts/test/TestProtocolFees.sol

@@ -21,11 +21,11 @@ pragma experimental ABIEncoderV2;
 
 import "@0x/contracts-asset-proxy/contracts/src/interfaces/IAssetProxy.sol";
 import "../src/interfaces/IStructs.sol";
-import "../src/Staking.sol";
+import "./TestStakingNoWETH.sol";
 
 
 contract TestProtocolFees is
-    Staking
+    TestStakingNoWETH
 {
     struct TestPool {
         uint256 operatorStake;
@@ -33,13 +33,22 @@ contract TestProtocolFees is
         mapping(address => bool) isMaker;
     }
 
+    event ERC20ProxyTransferFrom(
+        bytes assetData,
+        address from,
+        address to,
+        uint256 amount
+    );
+
     mapping(bytes32 => TestPool) private _testPools;
     mapping(address => bytes32) private _makersToTestPoolIds;
 
-    constructor(address exchangeAddress, address wethProxyAddress) public {
+    constructor(address exchangeAddress) public {
         init(
-            wethProxyAddress,
+            // Use this contract as the ERC20Proxy.
-            address(1), // vault addresses must be non-zero
+            address(this),
+            // vault addresses must be non-zero
+            address(1),
             address(1),
             address(1)
         );
@@ -81,6 +90,18 @@ contract TestProtocolFees is
         }
     }
 
+    /// @dev The ERC20Proxy `transferFrom()` function.
+    function transferFrom(
+        bytes calldata assetData,
+        address from,
+        address to,
+        uint256 amount
+    )
+        external
+    {
+        emit ERC20ProxyTransferFrom(assetData, from, to, amount);
+    }
+
     /// @dev Overridden to use test pools.
     function getStakingPoolIdOfMaker(address makerAddress)
         public

contracts/staking/contracts/test/TestStaking.sol

@@ -25,23 +25,17 @@ import "../src/Staking.sol";
 contract TestStaking is
     Staking
 {
-    address internal _wethAddress;
+    address public testWethAddress;
 
     constructor(address wethAddress) public {
-        _wethAddress = wethAddress;
+        testWethAddress = wethAddress;
     }
 
-    /// @dev Overridden to avoid hard-coded WETH.
+    /// @dev Overridden to use testWethAddress;
-    function getTotalBalance()
-        external
-        view
-        returns (uint256 totalBalance)
-    {
-        totalBalance = address(this).balance;
-    }
-
-    /// @dev Overridden to use _wethAddress;
     function _getWETHAddress() internal view returns (address) {
-        return _wethAddress;
+        // `testWethAddress` will not be set on the proxy this contract is
+        // attached to, so we need to access the storage of the deployed
+        // instance of this contract.
+        return TestStaking(address(uint160(stakingContract))).testWethAddress();
     }
 }

contracts/staking/contracts/test/TestStakingNoWETH.sol

@@ -17,28 +17,28 @@
 */
 
 pragma solidity ^0.5.9;
+pragma experimental ABIEncoderV2;
 
-import "@0x/contracts-asset-proxy/contracts/src/ERC20Proxy.sol";
+import "../src/Staking.sol";
 
 
-contract TestProtocolFeesERC20Proxy is
+// solhint-disable no-empty-blocks
-    ERC20Proxy
+/// @dev A version of the staking contract with WETH-related functions
+///      overridden to do nothing.
+contract TestStakingNoWETH is
+    Staking
 {
-    event TransferFromCalled(
+    function _transferWETHAllownces(
-        bytes assetData,
+        address[2] memory oldSpenders,
-        address from,
+        address[2] memory newSpenders
-        address to,
-        uint256 amount
-    );
-
-    function transferFrom(
-        bytes calldata assetData,
-        address from,
-        address to,
-        uint256 amount
     )
-        external
+        internal
+    {}
+
+    function _wrapBalanceToWETHAndGetBalance()
+        internal
+        returns (uint256 balance)
     {
-        emit TransferFromCalled(assetData, from, to, amount);
+        return address(this).balance;
     }
 }

contracts/staking/test/cumulative_reward_tracking_test.ts

@@ -2,8 +2,6 @@ import { ERC20Wrapper } from '@0x/contracts-asset-proxy';
 import { blockchainTests, describe } from '@0x/contracts-test-utils';
 import * as _ from 'lodash';
 
-import { artifacts } from '../src';
-
 import { deployAndConfigureContractsAsync, StakingApiWrapper } from './utils/api_wrapper';
 import { CumulativeRewardTrackingSimulation, TestAction } from './utils/cumulative_reward_tracking_simulation';
 

contracts/staking/test/params.ts

@@ -1,20 +1,28 @@
 import { blockchainTests, expect, filterLogsToArguments } from '@0x/contracts-test-utils';
 import { AuthorizableRevertErrors, BigNumber } from '@0x/utils';
+import { TransactionReceiptWithDecodedLogs } from 'ethereum-types';
+import * as _ from 'lodash';
 
-import { artifacts, IStakingEventsParamsSetEventArgs, MixinParamsContract } from '../src/';
+import {
+    artifacts,
+    IStakingEventsParamsSetEventArgs,
+    TestMixinParamsContract,
+    TestMixinParamsEvents,
+    TestMixinParamsWETHApproveEventArgs,
+} from '../src/';
 
 import { constants as stakingConstants } from './utils/constants';
 import { StakingParams } from './utils/types';
 
 blockchainTests('Configurable Parameters unit tests', env => {
-    let testContract: MixinParamsContract;
+    let testContract: TestMixinParamsContract;
     let authorizedAddress: string;
     let notAuthorizedAddress: string;
 
     before(async () => {
         [authorizedAddress, notAuthorizedAddress] = await env.getAccountAddressesAsync();
-        testContract = await MixinParamsContract.deployFrom0xArtifactAsync(
+        testContract = await TestMixinParamsContract.deployFrom0xArtifactAsync(
-            artifacts.MixinParams,
+            artifacts.TestMixinParams,
             env.provider,
             env.txDefaults,
             artifacts,
@@ -22,7 +30,7 @@ blockchainTests('Configurable Parameters unit tests', env => {
     });
 
     blockchainTests.resets('setParams()', () => {
-        async function setParamsAndAssertAsync(params: Partial<StakingParams>, from?: string): Promise<void> {
+        async function setParamsAndAssertAsync(params: Partial<StakingParams>, from?: string): Promise<TransactionReceiptWithDecodedLogs> {
             const _params = {
                 ...stakingConstants.DEFAULT_PARAMS,
                 ...params,
@@ -41,8 +49,9 @@ blockchainTests('Configurable Parameters unit tests', env => {
                 { from },
             );
             // Assert event.
-            expect(receipt.logs.length).to.eq(1);
+            const events = filterLogsToArguments<IStakingEventsParamsSetEventArgs>(receipt.logs, 'ParamsSet');
-            const event = filterLogsToArguments<IStakingEventsParamsSetEventArgs>(receipt.logs, 'ParamsSet')[0];
+            expect(events.length).to.eq(1);
+            const event = events[0];
             expect(event.epochDurationInSeconds).to.bignumber.eq(_params.epochDurationInSeconds);
             expect(event.rewardDelegatedStakeWeight).to.bignumber.eq(_params.rewardDelegatedStakeWeight);
             expect(event.minimumPoolStake).to.bignumber.eq(_params.minimumPoolStake);
@@ -65,6 +74,7 @@ blockchainTests('Configurable Parameters unit tests', env => {
             expect(actual[7]).to.eq(_params.ethVaultAddress);
             expect(actual[8]).to.eq(_params.rewardVaultAddress);
             expect(actual[9]).to.eq(_params.zrxVaultAddress);
+            return receipt;
         }
 
         it('throws if not called by an authorized address', async () => {

contracts/staking/test/protocol_fees.ts

@@ -17,8 +17,8 @@ import {
     IStakingEventsEvents,
     IStakingEventsStakingPoolActivatedEventArgs,
     TestProtocolFeesContract,
-    TestProtocolFeesERC20ProxyContract,
+    TestProtocolFeesERC20ProxyTransferFromEventArgs,
-    TestProtocolFeesERC20ProxyTransferFromCalledEventArgs,
+    TestProtocolFeesEvents,
 } from '../src';
 
 import { getRandomInteger } from './utils/number_utils';
@@ -34,14 +34,6 @@ blockchainTests('Protocol Fee Unit Tests', env => {
     before(async () => {
         [ownerAddress, exchangeAddress, notExchangeAddress] = await env.web3Wrapper.getAvailableAddressesAsync();
 
-        // Deploy the erc20Proxy for testing.
-        const proxy = await TestProtocolFeesERC20ProxyContract.deployFrom0xArtifactAsync(
-            artifacts.TestProtocolFeesERC20Proxy,
-            env.provider,
-            env.txDefaults,
-            {},
-        );
-
         // Deploy the protocol fees contract.
         testContract = await TestProtocolFeesContract.deployFrom0xArtifactAsync(
             artifacts.TestProtocolFees,
@@ -52,7 +44,6 @@ blockchainTests('Protocol Fee Unit Tests', env => {
             },
             artifacts,
             exchangeAddress,
-            proxy.address,
         );
 
         wethAssetData = await testContract.getWethAssetData.callAsync();
@@ -168,9 +159,9 @@ blockchainTests('Protocol Fee Unit Tests', env => {
 
         describe('ETH fees', () => {
             function assertNoWETHTransferLogs(logs: LogEntry[]): void {
-                const logsArgs = filterLogsToArguments<TestProtocolFeesERC20ProxyTransferFromCalledEventArgs>(
+                const logsArgs = filterLogsToArguments<TestProtocolFeesERC20ProxyTransferFromEventArgs>(
                     logs,
-                    'TransferFromCalled',
+                    TestProtocolFeesEvents.ERC20ProxyTransferFrom,
                 );
                 expect(logsArgs).to.deep.eq([]);
             }
@@ -233,9 +224,9 @@ blockchainTests('Protocol Fee Unit Tests', env => {
 
         describe('WETH fees', () => {
             function assertWETHTransferLogs(logs: LogEntry[], fromAddress: string, amount: BigNumber): void {
-                const logsArgs = filterLogsToArguments<TestProtocolFeesERC20ProxyTransferFromCalledEventArgs>(
+                const logsArgs = filterLogsToArguments<TestProtocolFeesERC20ProxyTransferFromEventArgs>(
                     logs,
-                    'TransferFromCalled',
+                    TestProtocolFeesEvents.ERC20ProxyTransferFrom,
                 );
                 expect(logsArgs.length).to.eq(1);
                 for (const args of logsArgs) {

contracts/staking/test/unit_tests/finalizer_test.ts

@@ -91,36 +91,35 @@ blockchainTests.resets('finalizer unit tests', env => {
     }
 
     interface FinalizationState {
-        balance: Numberish;
+        rewardsAvailable: Numberish;
-        currentEpoch: number;
+        poolsRemaining: number;
-        closingEpoch: number;
+        totalFeesCollected: Numberish;
-        numActivePoolsThisEpoch: number;
+        totalWeightedStake: Numberish;
-        totalFeesCollectedThisEpoch: Numberish;
+        totalRewardsFinalized: Numberish;
-        totalWeightedStakeThisEpoch: Numberish;
-        unfinalizedPoolsRemaining: number;
-        unfinalizedRewardsAvailable: Numberish;
-        unfinalizedTotalFeesCollected: Numberish;
-        unfinalizedTotalWeightedStake: Numberish;
     }
 
-    async function getFinalizationStateAsync(): Promise<FinalizationState> {
+    async function getUnfinalizedStateAsync(): Promise<FinalizationState> {
-        const r = await testContract.getFinalizationState.callAsync();
+        const r = await testContract.unfinalizedState.callAsync();
         return {
-            balance: r[0],
+            rewardsAvailable: r[0],
-            currentEpoch: r[1].toNumber(),
+            poolsRemaining: r[1].toNumber(),
-            closingEpoch: r[2].toNumber(),
+            totalFeesCollected: r[2],
-            numActivePoolsThisEpoch: r[3].toNumber(),
+            totalWeightedStake: r[3],
-            totalFeesCollectedThisEpoch: r[4],
+            totalRewardsFinalized: r[4],
-            totalWeightedStakeThisEpoch: r[5],
-            unfinalizedPoolsRemaining: r[6].toNumber(),
-            unfinalizedRewardsAvailable: r[7],
-            unfinalizedTotalFeesCollected: r[8],
-            unfinalizedTotalWeightedStake: r[9],
         };
     }
 
+    async function finalizePoolsAsync(poolIds: string[]): Promise<LogEntry[]> {
+        const logs = [] as LogEntry[];
+        for (const poolId of poolIds) {
+            const receipt = await testContract.finalizePool.awaitTransactionSuccessAsync(poolId);
+            logs.splice(logs.length - 1, 0, ...receipt.logs);
+        }
+        return logs;
+    }
+
     async function assertFinalizationStateAsync(expected: Partial<FinalizationState>): Promise<void> {
-        const actual = await getFinalizationStateAsync();
+        const actual = await getUnfinalizedStateAsync();
         assertEqualNumberFields(actual, expected);
     }
 
@@ -248,7 +247,7 @@ blockchainTests.resets('finalizer unit tests', env => {
         if (new BigNumber(pool.membersStake).isZero()) {
             return [new BigNumber(totalReward), ZERO_AMOUNT];
         }
-        const operatorShare = new BigNumber(totalReward).times(pool.operatorShare).integerValue(BigNumber.ROUND_DOWN);
+        const operatorShare = new BigNumber(totalReward).times(pool.operatorShare).integerValue(BigNumber.ROUND_UP);
         const membersShare = new BigNumber(totalReward).minus(operatorShare);
         return [operatorShare, membersShare];
     }
@@ -337,12 +336,12 @@ blockchainTests.resets('finalizer unit tests', env => {
             // Add a pool so there is state to clear.
             await addActivePoolAsync();
             await testContract.endEpoch.awaitTransactionSuccessAsync();
+            const epoch = await testContract.currentEpoch.callAsync();
+            expect(epoch).to.bignumber.eq(INITIAL_EPOCH + 1);
             return assertFinalizationStateAsync({
-                currentEpoch: INITIAL_EPOCH + 1,
+                poolsRemaining: 0,
-                closingEpoch: INITIAL_EPOCH,
+                totalFeesCollected: 0,
-                numActivePoolsThisEpoch: 0,
+                totalWeightedStake: 0,
-                totalFeesCollectedThisEpoch: 0,
-                totalWeightedStakeThisEpoch: 0,
             });
         });
 
@@ -351,10 +350,10 @@ blockchainTests.resets('finalizer unit tests', env => {
             const pool = await addActivePoolAsync();
             await testContract.endEpoch.awaitTransactionSuccessAsync();
             return assertFinalizationStateAsync({
-                unfinalizedPoolsRemaining: 1,
+                poolsRemaining: 1,
-                unfinalizedRewardsAvailable: INITIAL_BALANCE,
+                rewardsAvailable: INITIAL_BALANCE,
-                unfinalizedTotalFeesCollected: pool.feesCollected,
+                totalFeesCollected: pool.feesCollected,
-                unfinalizedTotalWeightedStake: pool.weightedStake,
+                totalWeightedStake: pool.weightedStake,
             });
         });
 
@@ -367,181 +366,35 @@ blockchainTests.resets('finalizer unit tests', env => {
         });
     });
 
-    describe('finalizePools()', () => {
-        it('does nothing if there were no active pools', async () => {
-            await testContract.endEpoch.awaitTransactionSuccessAsync();
-            const poolId = hexRandom();
-            const receipt = await testContract.finalizePools.awaitTransactionSuccessAsync([poolId]);
-            expect(receipt.logs).to.deep.eq([]);
-        });
-
-        it('does nothing if no pools are passed in', async () => {
-            await testContract.endEpoch.awaitTransactionSuccessAsync();
-            const receipt = await testContract.finalizePools.awaitTransactionSuccessAsync([]);
-            expect(receipt.logs).to.deep.eq([]);
-        });
-
-        it('can finalize a single pool', async () => {
-            const pool = await addActivePoolAsync();
-            await testContract.endEpoch.awaitTransactionSuccessAsync();
-            const receipt = await testContract.finalizePools.awaitTransactionSuccessAsync([pool.poolId]);
-            return assertFinalizationLogsAndBalancesAsync(INITIAL_BALANCE, [pool], receipt.logs);
-        });
-
-        it('can finalize multiple pools', async () => {
-            const pools = await Promise.all(_.times(3, async () => addActivePoolAsync()));
-            const poolIds = pools.map(p => p.poolId);
-            await testContract.endEpoch.awaitTransactionSuccessAsync();
-            const receipt = await testContract.finalizePools.awaitTransactionSuccessAsync(poolIds);
-            return assertFinalizationLogsAndBalancesAsync(INITIAL_BALANCE, pools, receipt.logs);
-        });
-
-        it('can finalize multiple pools over multiple transactions', async () => {
-            const pools = await Promise.all(_.times(2, async () => addActivePoolAsync()));
-            await testContract.endEpoch.awaitTransactionSuccessAsync();
-            const receipts = await Promise.all(
-                pools.map(pool => testContract.finalizePools.awaitTransactionSuccessAsync([pool.poolId])),
-            );
-            const allLogs = _.flatten(receipts.map(r => r.logs));
-            return assertFinalizationLogsAndBalancesAsync(INITIAL_BALANCE, pools, allLogs);
-        });
-
-        it('can finalize with no rewards', async () => {
-            await testContract.drainBalance.awaitTransactionSuccessAsync();
-            const pools = await Promise.all(_.times(2, async () => addActivePoolAsync()));
-            await testContract.endEpoch.awaitTransactionSuccessAsync();
-            const receipts = await Promise.all(
-                pools.map(pool => testContract.finalizePools.awaitTransactionSuccessAsync([pool.poolId])),
-            );
-            const allLogs = _.flatten(receipts.map(r => r.logs));
-            return assertFinalizationLogsAndBalancesAsync(0, pools, allLogs);
-        });
-
-        it('ignores a non-active pool', async () => {
-            const pools = await Promise.all(_.times(3, async () => addActivePoolAsync()));
-            const nonActivePoolId = hexRandom();
-            const poolIds = _.shuffle([...pools.map(p => p.poolId), nonActivePoolId]);
-            await testContract.endEpoch.awaitTransactionSuccessAsync();
-            const receipt = await testContract.finalizePools.awaitTransactionSuccessAsync(poolIds);
-            const rewardsPaidEvents = getRewardsPaidEvents(receipt.logs);
-            expect(rewardsPaidEvents.length).to.eq(pools.length);
-            for (const event of rewardsPaidEvents) {
-                expect(event.poolId).to.not.eq(nonActivePoolId);
-            }
-        });
-
-        it('ignores a finalized pool', async () => {
-            const pools = await Promise.all(_.times(3, async () => addActivePoolAsync()));
-            const poolIds = pools.map(p => p.poolId);
-            await testContract.endEpoch.awaitTransactionSuccessAsync();
-            const [finalizedPool] = _.sampleSize(pools, 1);
-            await testContract.finalizePools.awaitTransactionSuccessAsync([finalizedPool.poolId]);
-            const receipt = await testContract.finalizePools.awaitTransactionSuccessAsync(poolIds);
-            const rewardsPaidEvents = getRewardsPaidEvents(receipt.logs);
-            expect(rewardsPaidEvents.length).to.eq(pools.length - 1);
-            for (const event of rewardsPaidEvents) {
-                expect(event.poolId).to.not.eq(finalizedPool.poolId);
-            }
-        });
-
-        it('resets pool state after finalizing it', async () => {
-            const pools = await Promise.all(_.times(3, async () => addActivePoolAsync()));
-            const pool = _.sample(pools) as ActivePoolOpts;
-            await testContract.endEpoch.awaitTransactionSuccessAsync();
-            await testContract.finalizePools.awaitTransactionSuccessAsync([pool.poolId]);
-            const poolState = await testContract.getActivePoolFromEpoch.callAsync(
-                new BigNumber(INITIAL_EPOCH),
-                pool.poolId,
-            );
-            expect(poolState.feesCollected).to.bignumber.eq(0);
-            expect(poolState.weightedStake).to.bignumber.eq(0);
-            expect(poolState.membersStake).to.bignumber.eq(0);
-        });
-
-        it('`rewardsPaid` is the sum of all pool rewards', async () => {
-            const pools = await Promise.all(_.times(3, async () => addActivePoolAsync()));
-            const poolIds = pools.map(p => p.poolId);
-            await testContract.endEpoch.awaitTransactionSuccessAsync();
-            const receipt = await testContract.finalizePools.awaitTransactionSuccessAsync(poolIds);
-            const rewardsPaidEvents = getRewardsPaidEvents(receipt.logs);
-            const expectedTotalRewardsPaid = BigNumber.sum(
-                ...rewardsPaidEvents.map(e => e.membersReward.plus(e.operatorReward)),
-            );
-            const { rewardsPaid: totalRewardsPaid } = getEpochFinalizedEvents(receipt.logs)[0];
-            expect(totalRewardsPaid).to.bignumber.eq(expectedTotalRewardsPaid);
-        });
-
-        it('`rewardsPaid` <= `rewardsAvailable` <= contract balance at the end of the epoch', async () => {
-            const pools = await Promise.all(_.times(3, async () => addActivePoolAsync()));
-            const poolIds = pools.map(p => p.poolId);
-            let receipt = await testContract.endEpoch.awaitTransactionSuccessAsync();
-            const { rewardsAvailable } = getEpochEndedEvents(receipt.logs)[0];
-            expect(rewardsAvailable).to.bignumber.lte(INITIAL_BALANCE);
-            receipt = await testContract.finalizePools.awaitTransactionSuccessAsync(poolIds);
-            const { rewardsPaid } = getEpochFinalizedEvents(receipt.logs)[0];
-            expect(rewardsPaid).to.bignumber.lte(rewardsAvailable);
-        });
-
-        it('`rewardsPaid` <= `rewardsAvailable` with two equal pools', async () => {
-            const pool1 = await addActivePoolAsync();
-            const pool2 = await addActivePoolAsync(_.omit(pool1, 'poolId'));
-            const poolIds = [pool1, pool2].map(p => p.poolId);
-            let receipt = await testContract.endEpoch.awaitTransactionSuccessAsync();
-            const { rewardsAvailable } = getEpochEndedEvents(receipt.logs)[0];
-            receipt = await testContract.finalizePools.awaitTransactionSuccessAsync(poolIds);
-            const { rewardsPaid } = getEpochFinalizedEvents(receipt.logs)[0];
-            expect(rewardsPaid).to.bignumber.lte(rewardsAvailable);
-        });
-
-        blockchainTests.optional('`rewardsPaid` fuzzing', async () => {
-            const numTests = 32;
-            for (const i of _.times(numTests)) {
-                const numPools = _.random(1, 32);
-                it(`${i + 1}/${numTests} \`rewardsPaid\` <= \`rewardsAvailable\` (${numPools} pools)`, async () => {
-                    const pools = await Promise.all(_.times(numPools, async () => addActivePoolAsync()));
-                    const poolIds = pools.map(p => p.poolId);
-                    let receipt = await testContract.endEpoch.awaitTransactionSuccessAsync();
-                    const { rewardsAvailable } = getEpochEndedEvents(receipt.logs)[0];
-                    receipt = await testContract.finalizePools.awaitTransactionSuccessAsync(poolIds);
-                    const { rewardsPaid } = getEpochFinalizedEvents(receipt.logs)[0];
-                    expect(rewardsPaid).to.bignumber.lte(rewardsAvailable);
-                });
-            }
-        });
-    });
-
     describe('_finalizePool()', () => {
         it('does nothing if there were no active pools', async () => {
             await testContract.endEpoch.awaitTransactionSuccessAsync();
             const poolId = hexRandom();
-            const receipt = await testContract.finalizePool.awaitTransactionSuccessAsync(poolId);
+            const logs = await finalizePoolsAsync([poolId]);
-            expect(receipt.logs).to.deep.eq([]);
+            expect(logs).to.deep.eq([]);
         });
 
         it('can finalize a pool', async () => {
             const pool = await addActivePoolAsync();
             await testContract.endEpoch.awaitTransactionSuccessAsync();
-            const receipt = await testContract.finalizePool.awaitTransactionSuccessAsync(pool.poolId);
+            const logs = await finalizePoolsAsync([pool.poolId]);
-            return assertFinalizationLogsAndBalancesAsync(INITIAL_BALANCE, [pool], receipt.logs);
+            return assertFinalizationLogsAndBalancesAsync(INITIAL_BALANCE, [pool], logs);
         });
 
         it('can finalize multiple pools over multiple transactions', async () => {
             const pools = await Promise.all(_.times(2, async () => addActivePoolAsync()));
             await testContract.endEpoch.awaitTransactionSuccessAsync();
-            const receipts = await Promise.all(
+            const logs = await finalizePoolsAsync(pools.map(p => p.poolId));
-                pools.map(pool => testContract.finalizePool.awaitTransactionSuccessAsync(pool.poolId)),
+            return assertFinalizationLogsAndBalancesAsync(INITIAL_BALANCE, pools, logs);
-            );
-            const allLogs = _.flatten(receipts.map(r => r.logs));
-            return assertFinalizationLogsAndBalancesAsync(INITIAL_BALANCE, pools, allLogs);
         });
 
         it('ignores a finalized pool', async () => {
             const pools = await Promise.all(_.times(3, async () => addActivePoolAsync()));
             await testContract.endEpoch.awaitTransactionSuccessAsync();
             const [finalizedPool] = _.sampleSize(pools, 1);
-            await testContract.finalizePool.awaitTransactionSuccessAsync(finalizedPool.poolId);
+            await finalizePoolsAsync([finalizedPool.poolId]);
-            const receipt = await testContract.finalizePool.awaitTransactionSuccessAsync(finalizedPool.poolId);
+            const logs = await finalizePoolsAsync([finalizedPool.poolId]);
-            const rewardsPaidEvents = getRewardsPaidEvents(receipt.logs);
+            const rewardsPaidEvents = getRewardsPaidEvents(logs);
             expect(rewardsPaidEvents).to.deep.eq([]);
         });
 
@@ -549,7 +402,7 @@ blockchainTests.resets('finalizer unit tests', env => {
             const pools = await Promise.all(_.times(3, async () => addActivePoolAsync()));
             const pool = _.sample(pools) as ActivePoolOpts;
             await testContract.endEpoch.awaitTransactionSuccessAsync();
-            await testContract.finalizePool.awaitTransactionSuccessAsync(pool.poolId);
+            await finalizePoolsAsync([pool.poolId]);
             const poolState = await testContract.getActivePoolFromEpoch.callAsync(
                 new BigNumber(INITIAL_EPOCH),
                 pool.poolId,
@@ -564,11 +417,8 @@ blockchainTests.resets('finalizer unit tests', env => {
             const receipt = await testContract.endEpoch.awaitTransactionSuccessAsync();
             const { rewardsAvailable } = getEpochEndedEvents(receipt.logs)[0];
             expect(rewardsAvailable).to.bignumber.lte(INITIAL_BALANCE);
-            const receipts = await Promise.all(
+            const logs = await finalizePoolsAsync(pools.map(r => r.poolId));
-                pools.map(p => testContract.finalizePool.awaitTransactionSuccessAsync(p.poolId)),
+            const { rewardsPaid } = getEpochFinalizedEvents(logs)[0];
-            );
-            const allLogs = _.flatten(receipts.map(r => r.logs));
-            const { rewardsPaid } = getEpochFinalizedEvents(allLogs)[0];
             expect(rewardsPaid).to.bignumber.lte(rewardsAvailable);
         });
 
@@ -577,11 +427,8 @@ blockchainTests.resets('finalizer unit tests', env => {
             const pool2 = await addActivePoolAsync(_.omit(pool1, 'poolId'));
             const receipt = await testContract.endEpoch.awaitTransactionSuccessAsync();
             const { rewardsAvailable } = getEpochEndedEvents(receipt.logs)[0];
-            const receipts = await Promise.all(
+            const logs = await finalizePoolsAsync([pool1, pool2].map(r => r.poolId));
-                [pool1, pool2].map(p => testContract.finalizePool.awaitTransactionSuccessAsync(p.poolId)),
+            const { rewardsPaid } = getEpochFinalizedEvents(logs)[0];
-            );
-            const allLogs = _.flatten(receipts.map(r => r.logs));
-            const { rewardsPaid } = getEpochFinalizedEvents(allLogs)[0];
             expect(rewardsPaid).to.bignumber.lte(rewardsAvailable);
         });
 
@@ -593,11 +440,8 @@ blockchainTests.resets('finalizer unit tests', env => {
                     const pools = await Promise.all(_.times(numPools, async () => addActivePoolAsync()));
                     const receipt = await testContract.endEpoch.awaitTransactionSuccessAsync();
                     const { rewardsAvailable } = getEpochEndedEvents(receipt.logs)[0];
-                    const receipts = await Promise.all(
+                    const logs = await finalizePoolsAsync(pools.map(r => r.poolId));
-                        pools.map(p => testContract.finalizePool.awaitTransactionSuccessAsync(p.poolId)),
+                    const { rewardsPaid } = getEpochFinalizedEvents(logs)[0];
-                    );
-                    const allLogs = _.flatten(receipts.map(r => r.logs));
-                    const { rewardsPaid } = getEpochFinalizedEvents(allLogs)[0];
                     expect(rewardsPaid).to.bignumber.lte(rewardsAvailable);
                 });
             }
@@ -608,7 +452,7 @@ blockchainTests.resets('finalizer unit tests', env => {
         it('can advance the epoch after the prior epoch is finalized', async () => {
             const pool = await addActivePoolAsync();
             await testContract.endEpoch.awaitTransactionSuccessAsync();
-            await testContract.finalizePools.awaitTransactionSuccessAsync([pool.poolId]);
+            await finalizePoolsAsync([pool.poolId]);
             await testContract.endEpoch.awaitTransactionSuccessAsync();
             return expect(getCurrentEpochAsync()).to.become(INITIAL_EPOCH + 2);
         });
@@ -616,25 +460,25 @@ blockchainTests.resets('finalizer unit tests', env => {
         it('does not reward a pool that was only active 2 epochs ago', async () => {
             const pool1 = await addActivePoolAsync();
             await testContract.endEpoch.awaitTransactionSuccessAsync();
-            await testContract.finalizePools.awaitTransactionSuccessAsync([pool1.poolId]);
+            await finalizePoolsAsync([pool1.poolId]);
             await addActivePoolAsync();
             await testContract.endEpoch.awaitTransactionSuccessAsync();
             expect(getCurrentEpochAsync()).to.become(INITIAL_EPOCH + 2);
-            const receipt = await testContract.finalizePools.awaitTransactionSuccessAsync([pool1.poolId]);
+            const logs = await finalizePoolsAsync([pool1.poolId]);
-            const rewardsPaidEvents = getRewardsPaidEvents(receipt.logs);
+            const rewardsPaidEvents = getRewardsPaidEvents(logs);
             expect(rewardsPaidEvents).to.deep.eq([]);
         });
 
         it('does not reward a pool that was only active 3 epochs ago', async () => {
             const pool1 = await addActivePoolAsync();
             await testContract.endEpoch.awaitTransactionSuccessAsync();
-            await testContract.finalizePools.awaitTransactionSuccessAsync([pool1.poolId]);
+            await finalizePoolsAsync([pool1.poolId]);
             await testContract.endEpoch.awaitTransactionSuccessAsync();
             await addActivePoolAsync();
             await testContract.endEpoch.awaitTransactionSuccessAsync();
             expect(getCurrentEpochAsync()).to.become(INITIAL_EPOCH + 3);
-            const receipt = await testContract.finalizePools.awaitTransactionSuccessAsync([pool1.poolId]);
+            const logs = await finalizePoolsAsync([pool1.poolId]);
-            const rewardsPaidEvents = getRewardsPaidEvents(receipt.logs);
+            const rewardsPaidEvents = getRewardsPaidEvents(logs);
             expect(rewardsPaidEvents).to.deep.eq([]);
         });
 
@@ -642,11 +486,11 @@ blockchainTests.resets('finalizer unit tests', env => {
             const poolIds = _.times(3, () => hexRandom());
             await Promise.all(poolIds.map(async id => addActivePoolAsync({ poolId: id })));
             await testContract.endEpoch.awaitTransactionSuccessAsync();
-            let receipt = await testContract.finalizePools.awaitTransactionSuccessAsync(poolIds);
+            const finalizeLogs = await finalizePoolsAsync(poolIds);
-            const { rewardsRemaining: rolledOverRewards } = getEpochFinalizedEvents(receipt.logs)[0];
+            const { rewardsRemaining: rolledOverRewards } = getEpochFinalizedEvents(finalizeLogs)[0];
             await Promise.all(poolIds.map(async id => addActivePoolAsync({ poolId: id })));
-            receipt = await testContract.endEpoch.awaitTransactionSuccessAsync();
+            const {logs: endEpochLogs } = await testContract.endEpoch.awaitTransactionSuccessAsync();
-            const { rewardsAvailable } = getEpochEndedEvents(receipt.logs)[0];
+            const { rewardsAvailable } = getEpochEndedEvents(endEpochLogs)[0];
             expect(rewardsAvailable).to.bignumber.eq(rolledOverRewards);
         });
     });
@@ -694,7 +538,7 @@ blockchainTests.resets('finalizer unit tests', env => {
         it('returns empty if pool was only active in the 2 epochs ago', async () => {
             const pool = await addActivePoolAsync();
             await testContract.endEpoch.awaitTransactionSuccessAsync();
-            await testContract.finalizePools.awaitTransactionSuccessAsync([pool.poolId]);
+            await finalizePoolsAsync([pool.poolId]);
             return assertUnfinalizedPoolRewardsAsync(pool.poolId, ZERO_REWARDS);
         });
 
@@ -702,7 +546,7 @@ blockchainTests.resets('finalizer unit tests', env => {
             const pools = await Promise.all(_.times(3, async () => addActivePoolAsync()));
             const [pool] = _.sampleSize(pools, 1);
             await testContract.endEpoch.awaitTransactionSuccessAsync();
-            await testContract.finalizePools.awaitTransactionSuccessAsync([pool.poolId]);
+            await finalizePoolsAsync([pool.poolId]);
             return assertUnfinalizedPoolRewardsAsync(pool.poolId, ZERO_REWARDS);
         });
 

contracts/staking/test/utils/api_wrapper.ts

@@ -3,7 +3,7 @@ import { artifacts as erc20Artifacts, DummyERC20TokenContract, WETH9Contract } f
 import { BlockchainTestsEnvironment, constants, filterLogsToArguments, txDefaults } from '@0x/contracts-test-utils';
 import { BigNumber, logUtils } from '@0x/utils';
 import { Web3Wrapper } from '@0x/web3-wrapper';
-import { BlockParamLiteral, ContractArtifact, DecodedLogArgs, TransactionReceiptWithDecodedLogs } from 'ethereum-types';
+import { BlockParamLiteral, ContractArtifact, TransactionReceiptWithDecodedLogs } from 'ethereum-types';
 import * as _ from 'lodash';
 
 import {
@@ -12,29 +12,29 @@ import {
     IStakingEventsEpochEndedEventArgs,
     IStakingEventsStakingPoolActivatedEventArgs,
     ReadOnlyProxyContract,
-    StakingContract,
-    StakingEvents,
     StakingPoolRewardVaultContract,
     StakingProxyContract,
     TestCobbDouglasContract,
     TestStakingContract,
+    TestStakingEvents,
     ZrxVaultContract,
 } from '../../src';
 
 import { constants as stakingConstants } from './constants';
-import { EndOfEpochInfo, StakingParams } from './types';
+import { DecodedLogs, EndOfEpochInfo, StakingParams } from './types';
 
 export class StakingApiWrapper {
     // The address of the real Staking.sol contract
     public stakingContractAddress: string;
     // The StakingProxy.sol contract wrapped as a StakingContract to borrow API
-    public stakingContract: StakingContract;
+    public stakingContract: TestStakingContract;
     // The StakingProxy.sol contract as a StakingProxyContract
     public stakingProxyContract: StakingProxyContract;
     public zrxVaultContract: ZrxVaultContract;
     public ethVaultContract: EthVaultContract;
     public rewardVaultContract: StakingPoolRewardVaultContract;
     public zrxTokenContract: DummyERC20TokenContract;
+    public wethContract: WETH9Contract;
     public cobbDouglasContract: TestCobbDouglasContract;
     public utils = {
         // Epoch Utils
@@ -49,17 +49,17 @@ export class StakingApiWrapper {
             await this._web3Wrapper.mineBlockAsync();
         },
 
-        skipToNextEpochAndFinalizeAsync: async (): Promise<DecodedLogArgs[]> => {
+        skipToNextEpochAndFinalizeAsync: async (): Promise<DecodedLogs> => {
             await this.utils.fastForwardToNextEpochAsync();
             const endOfEpochInfo = await this.utils.endEpochAsync();
             let totalGasUsed = 0;
-            const allLogs = [] as LogEntry[];
+            const allLogs = [] as DecodedLogs;
             for (const poolId of endOfEpochInfo.activePoolIds) {
                 const receipt = await this.stakingContract.finalizePool.awaitTransactionSuccessAsync(
                     poolId,
                 );
                 totalGasUsed += receipt.gasUsed;
-                allLogs.splice(allLogs.length, 0, receipt.logs);
+                allLogs.splice(allLogs.length, 0, ...(receipt.logs as DecodedLogs));
             }
             logUtils.log(`Finalization cost ${totalGasUsed} gas`);
             return allLogs;
@@ -70,7 +70,7 @@ export class StakingApiWrapper {
             const receipt = await this.stakingContract.endEpoch.awaitTransactionSuccessAsync();
             const [epochEndedEvent] = filterLogsToArguments<IStakingEventsEpochEndedEventArgs>(
                 receipt.logs,
-                StakingEvents.EpochEnded,
+                TestStakingEvents.EpochEnded,
             );
             return {
                 closingEpoch: epochEndedEvent.epoch,
@@ -85,11 +85,11 @@ export class StakingApiWrapper {
             const _epoch = epoch !== undefined ? epoch : await this.stakingContract.currentEpoch.callAsync();
             const events = filterLogsToArguments<IStakingEventsStakingPoolActivatedEventArgs>(
                 await this.stakingContract.getLogsAsync(
-                    StakingEvents.StakingPoolActivated,
+                    TestStakingEvents.StakingPoolActivated,
                     { fromBlock: BlockParamLiteral.Earliest, toBlock: BlockParamLiteral.Latest },
                     { epoch: new BigNumber(_epoch) },
                 ),
-                StakingEvents.StakingPoolActivated,
+                TestStakingEvents.StakingPoolActivated,
             );
             return events.map(e => e.poolId);
         },
@@ -177,11 +177,12 @@ export class StakingApiWrapper {
         env: BlockchainTestsEnvironment,
         ownerAddress: string,
         stakingProxyContract: StakingProxyContract,
-        stakingContract: StakingContract,
+        stakingContract: TestStakingContract,
         zrxVaultContract: ZrxVaultContract,
         ethVaultContract: EthVaultContract,
         rewardVaultContract: StakingPoolRewardVaultContract,
         zrxTokenContract: DummyERC20TokenContract,
+        wethContract: WETH9Contract,
         cobbDouglasContract: TestCobbDouglasContract,
     ) {
         this._web3Wrapper = env.web3Wrapper;
@@ -189,13 +190,14 @@ export class StakingApiWrapper {
         this.ethVaultContract = ethVaultContract;
         this.rewardVaultContract = rewardVaultContract;
         this.zrxTokenContract = zrxTokenContract;
+        this.wethContract = wethContract;
         this.cobbDouglasContract = cobbDouglasContract;
 
         this.stakingContractAddress = stakingContract.address;
         this.stakingProxyContract = stakingProxyContract;
         // disguise the staking proxy as a StakingContract
         const logDecoderDependencies = _.mapValues({ ...artifacts, ...erc20Artifacts }, v => v.compilerOutput.abi);
-        this.stakingContract = new StakingContract(
+        this.stakingContract = new TestStakingContract(
             stakingProxyContract.address,
             env.provider,
             {
@@ -256,6 +258,7 @@ export async function deployAndConfigureContractsAsync(
         env.provider,
         env.txDefaults,
         artifacts,
+        wethContract.address,
     );
     // deploy reward vault
     const rewardVaultContract = await StakingPoolRewardVaultContract.deployFrom0xArtifactAsync(
@@ -263,6 +266,7 @@ export async function deployAndConfigureContractsAsync(
         env.provider,
         env.txDefaults,
         artifacts,
+        wethContract.address,
     );
     // deploy zrx vault
     const zrxVaultContract = await ZrxVaultContract.deployFrom0xArtifactAsync(
@@ -311,6 +315,7 @@ export async function deployAndConfigureContractsAsync(
         ethVaultContract,
         rewardVaultContract,
         zrxTokenContract,
+        wethContract,
         cobbDouglasContract,
     );
 }

contracts/staking/test/utils/cumulative_reward_tracking_simulation.ts

@@ -3,11 +3,11 @@ import { BigNumber } from '@0x/utils';
 import { DecodedLogEntry, TransactionReceiptWithDecodedLogs } from 'ethereum-types';
 import * as _ from 'lodash';
 
-import { artifacts, TestCumulativeRewardTrackingContract, IStakingEvents } from '../../src';
+import { artifacts, TestCumulativeRewardTrackingContract, TestCumulativeRewardTrackingEvents } from '../../src';
 
 import { StakingApiWrapper } from './api_wrapper';
 import { toBaseUnitAmount } from './number_utils';
-import { StakeInfo, StakeStatus } from './types';
+import { DecodedLogs, StakeInfo, StakeStatus } from './types';
 
 export enum TestAction {
     Finalize,
@@ -33,22 +33,22 @@ export class CumulativeRewardTrackingSimulation {
     private _testCumulativeRewardTrackingContract?: TestCumulativeRewardTrackingContract;
     private _poolId: string;
 
-    private static _extractTestLogs(txReceiptLogs: DecodedLogArgs[]): TestLog[] {
+    private static _extractTestLogs(txReceiptLogs: DecodedLogs): TestLog[] {
         const logs = [];
         for (const log of txReceiptLogs) {
-            if (log.event === 'SetMostRecentCumulativeReward') {
+            if (log.event === TestCumulativeRewardTrackingEvents.SetMostRecentCumulativeReward) {
                 logs.push({
-                    event: 'SetMostRecentCumulativeReward',
+                    event: log.event,
                     epoch: log.args.epoch.toNumber(),
                 });
-            } else if (log.event === 'SetCumulativeReward') {
+            } else if (log.event === TestCumulativeRewardTrackingEvents.SetCumulativeReward) {
                 logs.push({
-                    event: 'SetCumulativeReward',
+                    event: log.event,
                     epoch: log.args.epoch.toNumber(),
                 });
-            } else if (log.event === 'UnsetCumulativeReward') {
+            } else if (log.event === TestCumulativeRewardTrackingEvents.UnsetCumulativeReward) {
                 logs.push({
-                    event: 'UnsetCumulativeReward',
+                    event: log.event,
                     epoch: log.args.epoch.toNumber(),
                 });
             }
@@ -56,7 +56,7 @@ export class CumulativeRewardTrackingSimulation {
         return logs;
     }
 
-    private static _assertTestLogs(expectedSequence: TestLog[], txReceiptLogs: DecodedLogArgs[]): void {
+    private static _assertTestLogs(expectedSequence: TestLog[], txReceiptLogs: DecodedLogs): void {
         const logs = CumulativeRewardTrackingSimulation._extractTestLogs(txReceiptLogs);
         expect(logs.length).to.be.equal(expectedSequence.length);
         for (let i = 0; i < expectedSequence.length; i++) {
@@ -90,6 +90,7 @@ export class CumulativeRewardTrackingSimulation {
             env.provider,
             txDefaults,
             artifacts,
+            this._stakingApiWrapper.wethContract.address,
         );
     }
 
@@ -117,11 +118,11 @@ export class CumulativeRewardTrackingSimulation {
         CumulativeRewardTrackingSimulation._assertTestLogs(expectedTestLogs, testLogs);
     }
 
-    private async _executeActionsAsync(actions: TestAction[]): Promise<Array<DecodedLogEntry<any>>> {
+    private async _executeActionsAsync(actions: TestAction[]): Promise<DecodedLogs> {
-        const combinedLogs = [] as Array<DecodedLogEntry<any>>;
+        const combinedLogs = [] as DecodedLogs;
         for (const action of actions) {
-            let receipt: TransactionReceiptWithDecodedLogs;
+            let receipt: TransactionReceiptWithDecodedLogs | undefined;
-            let logs = [] as DecodedLogEntry<any>;
+            let logs = [] as DecodedLogs;
             switch (action) {
                 case TestAction.Finalize:
                     logs = await this._stakingApiWrapper.utils.skipToNextEpochAndFinalizeAsync();
@@ -163,15 +164,18 @@ export class CumulativeRewardTrackingSimulation {
                         true,
                         { from: this._poolOperator },
                     );
-                    const createStakingPoolLog = logs[0];
+                    const createStakingPoolLog = receipt.logs[0];
                     // tslint:disable-next-line no-unnecessary-type-assertion
-                    this._poolId = (createStakingPoolLog as DecodedLogArgs).args.poolId;
+                    this._poolId = (createStakingPoolLog as DecodedLogEntry<any>).args.poolId;
                     break;
 
                 default:
                     throw new Error('Unrecognized test action');
             }
-            combinedLogs.splice(combinedLogs.length - 1, 0, logs);
+            if (receipt !== undefined) {
+                logs = receipt.logs as DecodedLogs;
+            }
+            combinedLogs.splice(combinedLogs.length - 1, 0, ...logs);
         }
         return combinedLogs;
     }

contracts/staking/test/utils/types.ts

@@ -1,5 +1,6 @@
 import { Numberish } from '@0x/contracts-test-utils';
 import { BigNumber } from '@0x/utils';
+import { DecodedLogArgs, LogWithDecodedArgs } from 'ethereum-types';
 
 import { constants } from './constants';
 
@@ -133,3 +134,5 @@ export interface OperatorByPoolId {
 export interface DelegatorsByPoolId {
     [key: string]: string[];
 }
+
+export type DecodedLogs = Array<LogWithDecodedArgs<DecodedLogArgs>>;