@0x/contracts-staking: Reduce code duplication in MixinFinalizer and add unit tests for it.

contracts/staking/contracts/src/staking_pools/MixinStakingPoolRewards.sol

@@ -63,13 +63,13 @@ contract MixinStakingPoolRewards is
         view
         returns (uint256 reward)
     {
-        IStructs.PoolRewards memory unfinalizedPoolReward =
+        IStructs.PoolRewards memory unfinalizedPoolRewards =
-            _getUnfinalizedPoolReward(poolId);
+            _getUnfinalizedPoolRewards(poolId);
         reward = _computeRewardBalanceOfDelegator(
             poolId,
             member,
-            unfinalizedPoolReward.membersReward,
+            unfinalizedPoolRewards.membersReward,
-            unfinalizedPoolReward.membersStake
+            unfinalizedPoolRewards.membersStake
         );
     }
 

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

@@ -30,7 +30,7 @@ contract MixinAbstract {
     ///      Does nothing if the pool is already finalized.
     /// @param poolId The pool's ID.
     /// @return rewards Amount of rewards for this pool.
-    function _getUnfinalizedPoolReward(bytes32 poolId)
+    function _getUnfinalizedPoolRewards(bytes32 poolId)
         internal
         view
         returns (IStructs.PoolRewards memory rewards);

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

@@ -154,28 +154,19 @@ contract MixinFinalizer is
                 continue;
             }
 
-            // Clear the pool state so we don't finalize it again, and to
+            // Clear the pool state so we don't finalize it again, and to recoup
-            // recoup some gas.
+            // some gas.
             delete activePools[poolId];
 
-            // Credit the pool with rewards.
-            // We will transfer the total rewards to the vault at the end.
             IStructs.PoolRewards memory poolRewards =
-                _creditRewardToPool(poolId, pool);
+                _finalizePool(epoch, poolId, pool, true);
+
             rewardsPaid = rewardsPaid.safeAdd(
                 poolRewards.operatorReward + poolRewards.membersReward
             );
 
             // Decrease the number of unfinalized pools left.
             poolsRemaining = poolsRemaining.safeSub(1);
-
-            // Emit an event.
-            emit RewardsPaid(
-                epoch,
-                poolId,
-                poolRewards.operatorReward,
-                poolRewards.membersReward
-            );
         }
 
         // Deposit all the rewards at once into the RewardVault.
@@ -216,53 +207,32 @@ contract MixinFinalizer is
         if (epoch == 0) {
             return rewards;
         }
-
+        rewards = _finalizePool(
-        // Get the active pool.
-        mapping (bytes32 => IStructs.ActivePool) storage activePools =
-            _getActivePoolsFromEpoch(epoch - 1);
-        IStructs.ActivePool memory pool = activePools[poolId];
-
-        // Ignore pools that weren't active.
-        if (pool.feesCollected == 0) {
-            return rewards;
-        }
-
-        // Clear the pool state so we don't finalize it again, and to recoup
-        // some gas.
-        delete activePools[poolId];
-
-        // Credit the pool with rewards.
-        // We will transfer the total rewards to the vault at the end.
-        rewards = _creditRewardToPool(poolId, pool);
-        uint256 totalReward = rewards.membersReward + rewards.operatorReward;
-        totalRewardsPaidLastEpoch =
-            totalRewardsPaidLastEpoch.safeAdd(totalReward);
-
-        // Decrease the number of unfinalized pools left.
-        uint256 poolsRemaining =
-            unfinalizedPoolsRemaining =
-            unfinalizedPoolsRemaining.safeSub(1);
-
-        // Emit an event.
-        emit RewardsPaid(
             epoch,
             poolId,
-            rewards.operatorReward,
+            _getActivePoolFromEpoch(epoch - 1, poolId),
-            rewards.membersReward
+            false
         );
+    }
 
-        // Deposit all the rewards at once into the RewardVault.
+    /// @dev Computes the reward owed to a pool during finalization.
-        _depositIntoStakingPoolRewardVault(totalReward);
+    ///      Does nothing if the pool is already finalized.
-
+    /// @param poolId The pool's ID.
-        // If there are no more unfinalized pools remaining, the epoch is
+    /// @return rewards Amount of rewards for this pool.
-        // finalized.
+    function _getUnfinalizedPoolRewards(bytes32 poolId)
-        if (poolsRemaining == 0) {
+        internal
-            emit EpochFinalized(
+        view
-                epoch - 1,
+        returns (IStructs.PoolRewards memory rewards)
-                totalRewardsPaidLastEpoch,
+    {
-                unfinalizedRewardsAvailable.safeSub(totalRewardsPaidLastEpoch)
+        uint256 epoch = getCurrentEpoch();
-            );
+        // There are no pools to finalize at epoch 0.
+        if (epoch == 0) {
+            return rewards;
         }
+        rewards = _getUnfinalizedPoolRewards(
+            poolId,
+            _getActivePoolFromEpoch(epoch - 1, poolId)
+        );
     }
 
     /// @dev Get an active pool from an epoch by its ID.
@@ -295,50 +265,6 @@ contract MixinFinalizer is
         activePools = activePoolsByEpoch[epoch % 2];
     }
 
-    /// @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 rewards Amount of rewards for this pool.
-    function _getUnfinalizedPoolReward(bytes32 poolId)
-        internal
-        view
-        returns (IStructs.PoolRewards memory rewards)
-    {
-        uint256 epoch = getCurrentEpoch();
-        // There can't be any rewards in the first epoch.
-        if (epoch == 0) {
-            return rewards;
-        }
-
-        IStructs.ActivePool memory pool =
-            _getActivePoolFromEpoch(epoch - 1, poolId);
-        // There can't be any rewards if the pool was active or if it has
-        // no stake.
-        if (pool.feesCollected == 0 || pool.weightedStake == 0) {
-            return rewards;
-        }
-
-        // Use the cobb-douglas function to compute the total reward.
-        uint256 totalReward = LibCobbDouglas._cobbDouglas(
-            unfinalizedRewardsAvailable,
-            pool.feesCollected,
-            unfinalizedTotalFeesCollected,
-            pool.weightedStake,
-            unfinalizedTotalWeightedStake,
-            cobbDouglasAlphaNumerator,
-            cobbDouglasAlphaDenomintor
-        );
-
-        // Split the reward between the operator and delegators.
-        if (pool.membersStake == 0) {
-            rewards.operatorReward = totalReward;
-        } else {
-            (rewards.operatorReward, rewards.membersReward) =
-                _splitAmountBetweenOperatorAndMembers(poolId, totalReward);
-        }
-        rewards.membersStake = pool.membersStake;
-    }
-
     /// @dev Converts the entire WETH balance of the contract into ETH.
     function _unwrapWETH() internal {
         uint256 wethBalance = IEtherToken(WETH_ADDRESS)
@@ -354,7 +280,7 @@ contract MixinFinalizer is
     /// @param amount Amount to to split.
     /// @return operatorPortion Portion of `amount` attributed to the operator.
     /// @return membersPortion Portion of `amount` attributed to the pool.
-    function _splitAmountBetweenOperatorAndMembers(
+    function _splitRewardAmountBetweenOperatorAndMembers(
         bytes32 poolId,
         uint256 amount
     )
@@ -390,21 +316,21 @@ contract MixinFinalizer is
         );
     }
 
-    /// @dev Computes the reward owed to a pool during finalization and
+    /// @dev Computes the reward owed to a pool during finalization.
-    ///      credits it to that pool for the CURRENT epoch.
     /// @param poolId The pool's ID.
-    /// @param pool The pool.
+    /// @param pool The active pool.
     /// @return rewards Amount of rewards for this pool.
-    function _creditRewardToPool(
+    function _getUnfinalizedPoolRewards(
         bytes32 poolId,
         IStructs.ActivePool memory pool
     )
         private
+        view
         returns (IStructs.PoolRewards memory rewards)
     {
         // There can't be any rewards if the pool was active or if it has
         // no stake.
-        if (pool.feesCollected == 0 || pool.weightedStake == 0) {
+        if (pool.feesCollected == 0) {
             return rewards;
         }
 
@@ -419,15 +345,59 @@ contract MixinFinalizer is
             cobbDouglasAlphaDenomintor
         );
 
-        // Credit the pool the reward in the RewardVault.
+        // Split the reward between the operator and delegators.
-        (rewards.operatorReward, rewards.membersReward) =
+        if (pool.membersStake == 0) {
-            _recordDepositInRewardVaultFor(
+            rewards.operatorReward = totalReward;
-                poolId,
+        } else {
-                totalReward,
+            (rewards.operatorReward, rewards.membersReward) =
-                // If no delegated stake, all rewards go to the operator.
+                _splitRewardAmountBetweenOperatorAndMembers(
-                pool.membersStake == 0
+                    poolId,
-            );
+                    totalReward
+                );
+        }
         rewards.membersStake = pool.membersStake;
+    }
+
+    /// @dev Either fully or partially finalizes a single pool that was active
+    ///      in the previous epoch. If `batchedMode` is `true`, this function
+    ///      will NOT:
+    ///         - transfer ether into the reward vault
+    ///         - update `poolsRemaining`
+    ///         - update `totalRewardsPaidLastEpoch`
+    ///         - clear the pool from `activePoolsByEpoch`
+    ///         - emit an `EpochFinalized` event.
+    /// @param epoch The current epoch.
+    /// @param poolId The pool ID to finalize.
+    /// @param pool The active pool to finalize.
+    /// @param batchedMode Only calculate and credit rewards.
+    /// @return rewards Rewards.
+    /// @return rewards The rewards credited to the pool.
+    function _finalizePool(
+        uint256 epoch,
+        bytes32 poolId,
+        IStructs.ActivePool memory pool,
+        bool batchedMode
+    )
+        private
+        returns (IStructs.PoolRewards memory rewards)
+    {
+        // Ignore pools that weren't active.
+        if (pool.feesCollected == 0) {
+            return rewards;
+        }
+
+        // Compute the rewards.
+        rewards = _getUnfinalizedPoolRewards(poolId, pool);
+        uint256 totalReward =
+            rewards.membersReward.safeAdd(rewards.operatorReward);
+
+        // Credit the pool the rewards in the RewardVault.
+        _recordDepositInRewardVaultFor(
+            poolId,
+            totalReward,
+            // If no delegated stake, all rewards go to the operator.
+            pool.membersStake == 0
+        );
 
         // Sync delegator rewards.
         if (rewards.membersReward != 0) {
@@ -437,5 +407,41 @@ contract MixinFinalizer is
                 pool.membersStake
             );
         }
+
+        // Emit an event.
+        emit RewardsPaid(
+            epoch,
+            poolId,
+            rewards.operatorReward,
+            rewards.membersReward
+        );
+
+        if (batchedMode) {
+            return rewards;
+        }
+
+        // Clear the pool state so we don't finalize it again, and to recoup
+        // some gas.
+        delete _getActivePoolsFromEpoch(epoch)[poolId];
+
+        if (totalReward > 0) {
+            totalRewardsPaidLastEpoch =
+                totalRewardsPaidLastEpoch.safeAdd(totalReward);
+            _depositIntoStakingPoolRewardVault(totalReward);
+        }
+
+        // Decrease the number of unfinalized pools left.
+        uint256 poolsRemaining = unfinalizedPoolsRemaining;
+        unfinalizedPoolsRemaining = poolsRemaining = poolsRemaining.safeSub(1);
+
+        // If there are no more unfinalized pools remaining, the epoch is
+        // finalized.
+        if (poolsRemaining == 0) {
+            emit EpochFinalized(
+                epoch - 1,
+                totalRewardsPaidLastEpoch,
+                unfinalizedRewardsAvailable.safeSub(totalRewardsPaidLastEpoch)
+            );
+        }
     }
 }

contracts/staking/contracts/test/TestFinalizer.sol

@@ -26,37 +26,58 @@ import "../src/Staking.sol";
 contract TestFinalizer is
     Staking
 {
-    struct RecordedReward {
+    event RecordRewardForDelegatorsCall(
-        uint256 membersReward;
+        bytes32 poolId,
-        uint256 membersStake;
+        uint256 membersReward,
-    }
+        uint256 membersStake
+    );
+
+    event RecordDepositInRewardVaultForCall(
+        bytes32 poolId,
+        uint256 totalReward,
+        bool operatorOnly
+    );
+
+    event DepositIntoStakingPoolRewardVaultCall(
+        uint256 amount
+    );
 
-    struct DepositedReward {
-        uint256 totalReward;
-        bool operatorOnly;
-    }
     mapping (bytes32 => uint32) internal _operatorSharesByPool;
-    mapping (bytes32 => RecordedReward) internal _recordedRewardsByPool;
-    mapping (bytes32 => DepositedReward) internal _depositedRewardsByPool;
 
+    constructor() public {
+        init();
+    }
+
+    /// @dev Get finalization-related state variables.
     function getFinalizationState()
         external
         view
         returns (
+            uint256 _balance,
+            uint256 _currentEpoch,
             uint256 _closingEpoch,
+            uint256 _numActivePoolsThisEpoch,
+            uint256 _totalFeesCollectedThisEpoch,
+            uint256 _totalWeightedStakeThisEpoch,
             uint256 _unfinalizedPoolsRemaining,
             uint256 _unfinalizedRewardsAvailable,
             uint256 _unfinalizedTotalFeesCollected,
             uint256 _unfinalizedTotalWeightedStake
         )
     {
+        _balance = address(this).balance;
+        _currentEpoch = currentEpoch;
         _closingEpoch = currentEpoch - 1;
+        _numActivePoolsThisEpoch = numActivePoolsThisEpoch;
+        _totalFeesCollectedThisEpoch = totalFeesCollectedThisEpoch;
+        _totalWeightedStakeThisEpoch = totalWeightedStakeThisEpoch;
         _unfinalizedPoolsRemaining = unfinalizedPoolsRemaining;
         _unfinalizedRewardsAvailable = unfinalizedRewardsAvailable;
         _unfinalizedTotalFeesCollected = unfinalizedTotalFeesCollected;
         _unfinalizedTotalWeightedStake = unfinalizedTotalWeightedStake;
     }
 
+    /// @dev Activate a pool in the current epoch.
     function addActivePool(
         bytes32 poolId,
         uint32 operatorShare,
@@ -66,9 +87,10 @@ contract TestFinalizer is
     )
         external
     {
+        require(feesCollected > 0, "FEES_MUST_BE_NONZERO");
         mapping (bytes32 => IStructs.ActivePool) storage activePools =
             _getActivePoolsFromEpoch(currentEpoch);
-        assert(activePools[poolId].feesCollected == 0);
+        require(feesCollected > 0, "POOL_ALREADY_ADDED");
         _operatorSharesByPool[poolId] = operatorShare;
         activePools[poolId] = IStructs.ActivePool({
             feesCollected: feesCollected,
@@ -80,6 +102,34 @@ contract TestFinalizer is
         numActivePoolsThisEpoch += 1;
     }
 
+    /// @dev Expose `_getUnfinalizedPoolReward()`
+    function internalGetUnfinalizedPoolRewards(bytes32 poolId)
+        external
+        view
+        returns (IStructs.PoolRewards memory rewards)
+    {
+        rewards = _getUnfinalizedPoolRewards(poolId);
+    }
+
+
+    /// @dev Expose `_getActivePoolFromEpoch`.
+    function internalGetActivePoolFromEpoch(uint256 epoch, bytes32 poolId)
+        external
+        view
+        returns (IStructs.ActivePool memory pool)
+    {
+        pool = _getActivePoolFromEpoch(epoch, poolId);
+    }
+
+
+    /// @dev Expose `_finalizePool()`
+    function internalFinalizePool(bytes32 poolId)
+        external
+        returns (IStructs.PoolRewards memory rewards)
+    {
+        rewards = _finalizePool(poolId);
+    }
+
     /// @dev Overridden to just store inputs.
     function _recordRewardForDelegators(
         bytes32 poolId,
@@ -88,10 +138,16 @@ contract TestFinalizer is
     )
         internal
     {
-        _recordedRewardsByPool[poolId] = RecordedReward({
+        emit RecordRewardForDelegatorsCall(
-            membersReward: membersReward,
+            poolId,
-            membersStake: membersStake
+            membersReward,
-        });
+            membersStake
+        );
+    }
+
+    /// @dev Overridden to store inputs and do some really basic math.
+    function _depositIntoStakingPoolRewardVault(uint256 amount) internal {
+        emit DepositIntoStakingPoolRewardVaultCall(amount);
     }
 
     /// @dev Overridden to store inputs and do some really basic math.
@@ -106,21 +162,25 @@ contract TestFinalizer is
             uint256 membersPortion
         )
     {
-        _depositedRewardsByPool[poolId] = DepositedReward({
+        emit RecordDepositInRewardVaultForCall(
-            totalReward: totalReward,
+            poolId,
-            operatorOnly: operatorOnly
+            totalReward,
-        });
+            operatorOnly
+        );
 
         if (operatorOnly) {
             operatorPortion = totalReward;
         } else {
             (operatorPortion, membersPortion) =
-                _splitAmountBetweenOperatorAndMembers(poolId, totalReward);
+                _splitRewardAmountBetweenOperatorAndMembers(
+                    poolId,
+                    totalReward
+                );
         }
     }
 
     /// @dev Overridden to do some really basic math.
-    function _splitAmountBetweenOperatorAndMembers(
+    function _splitRewardAmountBetweenOperatorAndMembers(
         bytes32 poolId,
         uint256 amount
     )
@@ -133,7 +193,7 @@ contract TestFinalizer is
         membersPortion = amount - operatorPortion;
     }
 
-    /// @dev Overriden to always succeed.
+    /// @dev Overriden to just increase the epoch counter.
     function _goToNextEpoch() internal {
         currentEpoch += 1;
     }

contracts/staking/test/unit_tests/finalizer_test.ts

@@ -1,55 +1,499 @@
-import { blockchainTests, expect, filterLogsToArguments, Numberish } from '@0x/contracts-test-utils';
+import {
+    blockchainTests,
+    constants,
+    expect,
+    filterLogsToArguments,
+    hexRandom,
+    Numberish,
+} from '@0x/contracts-test-utils';
+import { StakingRevertErrors } from '@0x/order-utils';
+import { BigNumber } from '@0x/utils';
+import { LogEntry } from 'ethereum-types';
+import * as _ from 'lodash';
 
 import {
     artifacts,
     IStakingEventsEpochEndedEventArgs,
     IStakingEventsEpochFinalizedEventArgs,
     IStakingEventsEvents,
+    IStakingEventsRewardsPaidEventArgs,
     TestFinalizerContract,
+    TestFinalizerDepositIntoStakingPoolRewardVaultCallEventArgs,
+    TestFinalizerEvents,
 } from '../../src';
+import { getRandomInteger, toBaseUnitAmount } from '../utils/number_utils';
 
 blockchainTests.resets.only('finalization tests', env => {
-    let testContract: TestFinalizerContract;
+    const { ONE_ETHER, ZERO_AMOUNT } = constants;
     const INITIAL_EPOCH = 0;
+    const INITIAL_BALANCE = toBaseUnitAmount(32);
+    let senderAddress: string;
+    let testContract: TestFinalizerContract;
 
     before(async () => {
+        [senderAddress] = await env.getAccountAddressesAsync();
         testContract = await TestFinalizerContract.deployFrom0xArtifactAsync(
             artifacts.TestFinalizer,
             env.provider,
             env.txDefaults,
             artifacts,
         );
+        // Give the contract a balance.
+        await sendEtherAsync(testContract.address, INITIAL_BALANCE);
     });
 
-    describe('endEpoch()', () => {
+    async function sendEtherAsync(to: string, amount: Numberish): Promise<void> {
-        it('emits an `EpochEnded` event', async () => {
+        await env.web3Wrapper.awaitTransactionSuccessAsync(
-            const receipt = await testContract.endEpoch.awaitTransactionSuccessAsync();
+                await env.web3Wrapper.sendTransactionAsync({
-            const [epochEndedEvent] = filterLogsToArguments<IStakingEventsEpochEndedEventArgs>(
+                from: senderAddress,
-                receipt.logs,
+                to,
-                IStakingEventsEvents.EpochEnded,
+                value: new BigNumber(amount),
-            );
+            }),
-            expect(epochEndedEvent.epoch).to.bignumber.eq(INITIAL_EPOCH);
+        );
-            expect(epochEndedEvent.numActivePools).to.bignumber.eq(0);
+    }
-            expect(epochEndedEvent.rewardsAvailable).to.bignumber.eq(0);
-            expect(epochEndedEvent.totalFeesCollected).to.bignumber.eq(0);
-            expect(epochEndedEvent.totalWeightedStake).to.bignumber.eq(0);
-        });
 
+    interface ActivePoolOpts {
+        poolId: string;
+        operatorShare: number;
+        feesCollected: Numberish;
+        membersStake: Numberish;
+        weightedStake: Numberish;
+    }
+
+    async function addActivePoolAsync(opts?: Partial<ActivePoolOpts>): Promise<ActivePoolOpts> {
+         const _opts = {
+             poolId: hexRandom(),
+             operatorShare: Math.random(),
+             feesCollected: getRandomInteger(0, ONE_ETHER),
+             membersStake: getRandomInteger(0, ONE_ETHER),
+             weightedStake: getRandomInteger(0, ONE_ETHER),
+             ...opts,
+         };
+         await testContract.addActivePool.awaitTransactionSuccessAsync(
+             _opts.poolId,
+            new BigNumber(_opts.operatorShare * constants.PPM_DENOMINATOR).integerValue(),
+            new BigNumber(_opts.feesCollected),
+            new BigNumber(_opts.membersStake),
+            new BigNumber(_opts.weightedStake),
+         );
+         return _opts;
+    }
+
+    interface FinalizationState {
+        balance: Numberish;
+        currentEpoch: number;
+        closingEpoch: number;
+        numActivePoolsThisEpoch: number;
+        totalFeesCollectedThisEpoch: Numberish;
+        totalWeightedStakeThisEpoch: Numberish;
+        unfinalizedPoolsRemaining: number;
+        unfinalizedRewardsAvailable: Numberish;
+        unfinalizedTotalFeesCollected: Numberish;
+        unfinalizedTotalWeightedStake: Numberish;
+    }
+
+    async function getFinalizationStateAsync(): Promise<FinalizationState> {
+        const r = await testContract.getFinalizationState.callAsync();
+        return {
+            balance: r[0],
+            currentEpoch: r[1].toNumber(),
+            closingEpoch: r[2].toNumber(),
+            numActivePoolsThisEpoch: r[3].toNumber(),
+            totalFeesCollectedThisEpoch: r[4],
+            totalWeightedStakeThisEpoch: r[5],
+            unfinalizedPoolsRemaining: r[6].toNumber(),
+            unfinalizedRewardsAvailable: r[7],
+            unfinalizedTotalFeesCollected: r[8],
+            unfinalizedTotalWeightedStake: r[9],
+        };
+    }
+
+    async function assertFinalizationStateAsync(
+        expected: Partial<FinalizationState>,
+    ): Promise<void> {
+        const actual = await getFinalizationStateAsync();
+        if (expected.balance !== undefined) {
+            expect(actual.balance).to.bignumber.eq(expected.balance);
+        }
+        if (expected.currentEpoch !== undefined) {
+            expect(actual.currentEpoch).to.eq(expected.currentEpoch);
+        }
+        if (expected.closingEpoch !== undefined) {
+            expect(actual.closingEpoch).to.eq(expected.closingEpoch);
+        }
+        if (expected.numActivePoolsThisEpoch !== undefined) {
+            expect(actual.numActivePoolsThisEpoch)
+                .to.eq(expected.numActivePoolsThisEpoch);
+        }
+        if (expected.totalFeesCollectedThisEpoch !== undefined) {
+            expect(actual.totalFeesCollectedThisEpoch)
+                .to.bignumber.eq(expected.totalFeesCollectedThisEpoch);
+        }
+        if (expected.totalWeightedStakeThisEpoch !== undefined) {
+            expect(actual.totalWeightedStakeThisEpoch)
+                .to.bignumber.eq(expected.totalWeightedStakeThisEpoch);
+        }
+        if (expected.unfinalizedPoolsRemaining !== undefined) {
+            expect(actual.unfinalizedPoolsRemaining)
+                .to.eq(expected.unfinalizedPoolsRemaining);
+        }
+        if (expected.unfinalizedRewardsAvailable !== undefined) {
+            expect(actual.unfinalizedRewardsAvailable)
+                .to.bignumber.eq(expected.unfinalizedRewardsAvailable);
+        }
+        if (expected.unfinalizedTotalFeesCollected !== undefined) {
+            expect(actual.unfinalizedTotalFeesCollected)
+                .to.bignumber.eq(expected.unfinalizedTotalFeesCollected);
+        }
+        if (expected.unfinalizedTotalFeesCollected !== undefined) {
+            expect(actual.unfinalizedTotalFeesCollected)
+                .to.bignumber.eq(expected.unfinalizedTotalFeesCollected);
+        }
+    }
+
+    function assertEpochEndedEvent(
+        logs: LogEntry[],
+        args: Partial<IStakingEventsEpochEndedEventArgs>,
+    ): void {
+        const events = filterLogsToArguments<IStakingEventsEpochEndedEventArgs>(
+            logs,
+            IStakingEventsEvents.EpochEnded,
+        );
+        expect(events.length).to.eq(1);
+        if (args.epoch !== undefined) {
+            expect(events[0].epoch).to.bignumber.eq(INITIAL_EPOCH);
+        }
+        if (args.numActivePools !== undefined) {
+            expect(events[0].numActivePools).to.bignumber.eq(args.numActivePools);
+        }
+        if (args.rewardsAvailable !== undefined) {
+            expect(events[0].rewardsAvailable).to.bignumber.eq(args.rewardsAvailable);
+        }
+        if (args.totalFeesCollected !== undefined) {
+            expect(events[0].totalFeesCollected).to.bignumber.eq(args.totalFeesCollected);
+        }
+        if (args.totalWeightedStake !== undefined) {
+            expect(events[0].totalWeightedStake).to.bignumber.eq(args.totalWeightedStake);
+        }
+    }
+
+    function assertEpochFinalizedEvent(
+            logs: LogEntry[],
+            args: Partial<IStakingEventsEpochFinalizedEventArgs>,
+    ): void {
+        const events = getEpochFinalizedEvents(logs);
+        expect(events.length).to.eq(1);
+        if (args.epoch !== undefined) {
+            expect(events[0].epoch).to.bignumber.eq(args.epoch);
+        }
+        if (args.rewardsPaid !== undefined) {
+            expect(events[0].rewardsPaid).to.bignumber.eq(args.rewardsPaid);
+        }
+        if (args.rewardsRemaining !== undefined) {
+            expect(events[0].rewardsRemaining).to.bignumber.eq(args.rewardsRemaining);
+        }
+    }
+
+    function assertDepositIntoStakingPoolRewardVaultCallEvent(
+        logs: LogEntry[],
+        amount?: Numberish,
+    ): void {
+        const events = filterLogsToArguments<TestFinalizerDepositIntoStakingPoolRewardVaultCallEventArgs>(
+            logs,
+            TestFinalizerEvents.DepositIntoStakingPoolRewardVaultCall,
+        );
+        expect(events.length).to.eq(1);
+        if (amount !== undefined) {
+            expect(events[0].amount).to.bignumber.eq(amount);
+        }
+    }
+
+    function getEpochFinalizedEvents(logs: LogEntry[]): IStakingEventsEpochFinalizedEventArgs[] {
+        return filterLogsToArguments<IStakingEventsEpochFinalizedEventArgs>(
+            logs,
+            IStakingEventsEvents.EpochFinalized,
+        );
+    }
+
+    function getRewardsPaidEvents(logs: LogEntry[]): IStakingEventsRewardsPaidEventArgs[] {
+        return filterLogsToArguments<IStakingEventsRewardsPaidEventArgs>(
+            logs,
+            IStakingEventsEvents.RewardsPaid,
+        );
+    }
+
+    async function getCurrentEpochAsync(): Promise<number> {
+        return (await testContract.getCurrentEpoch.callAsync()).toNumber();
+    }
+
+    describe('endEpoch()', () => {
         it('advances the epoch', async () => {
             await testContract.endEpoch.awaitTransactionSuccessAsync();
             const currentEpoch = await testContract.getCurrentEpoch.callAsync();
-            expect(currentEpoch).to.be.bignumber.eq(INITIAL_EPOCH + 1);
+            expect(currentEpoch).to.bignumber.eq(INITIAL_EPOCH + 1);
+        });
+
+        it('emits an `EpochEnded` event', async () => {
+            const receipt = await testContract.endEpoch.awaitTransactionSuccessAsync();
+            assertEpochEndedEvent(
+                receipt.logs,
+                {
+                    epoch: new BigNumber(INITIAL_EPOCH),
+                    numActivePools: ZERO_AMOUNT,
+                    rewardsAvailable: INITIAL_BALANCE,
+                    totalFeesCollected: ZERO_AMOUNT,
+                    totalWeightedStake: ZERO_AMOUNT,
+                },
+            );
         });
 
         it('immediately finalizes if there are no active pools', async () => {
             const receipt = await testContract.endEpoch.awaitTransactionSuccessAsync();
-            const [epochFinalizedEvent] = filterLogsToArguments<IStakingEventsEpochFinalizedEventArgs>(
+            assertEpochFinalizedEvent(
+                receipt.logs,
+                {
+                    epoch: new BigNumber(INITIAL_EPOCH),
+                    rewardsPaid: ZERO_AMOUNT,
+                    rewardsRemaining: INITIAL_BALANCE,
+                },
+            );
+        });
+
+        it('does not immediately finalize if there is an active pool', async () => {
+            await addActivePoolAsync();
+            const receipt = await testContract.endEpoch.awaitTransactionSuccessAsync();
+            const events = filterLogsToArguments<IStakingEventsEpochFinalizedEventArgs>(
                 receipt.logs,
                 IStakingEventsEvents.EpochFinalized,
             );
-            expect(epochFinalizedEvent.epoch).to.bignumber.eq(INITIAL_EPOCH);
+            expect(events).to.deep.eq([]);
-            expect(epochFinalizedEvent.rewardsPaid).to.bignumber.eq(0);
+        });
-            expect(epochFinalizedEvent.rewardsRemaining).to.bignumber.eq(0);
+
+        it('clears the next epoch\'s finalization state', async () => {
+            // Add a pool so there is state to clear.
+            await addActivePoolAsync();
+            await testContract.endEpoch.awaitTransactionSuccessAsync();
+            assertFinalizationStateAsync({
+                currentEpoch: INITIAL_EPOCH + 1,
+                closingEpoch: INITIAL_EPOCH,
+                numActivePoolsThisEpoch: 0,
+                totalFeesCollectedThisEpoch: 0,
+                totalWeightedStakeThisEpoch: 0,
+            });
+        });
+
+        it('prepares finalization state', async () => {
+            // Add a pool so there is state to clear.
+            const pool = await addActivePoolAsync();
+            await testContract.endEpoch.awaitTransactionSuccessAsync();
+            assertFinalizationStateAsync({
+                unfinalizedPoolsRemaining: 1,
+                unfinalizedRewardsAvailable: INITIAL_BALANCE,
+                unfinalizedTotalFeesCollected: pool.feesCollected,
+                unfinalizedTotalWeightedStake: pool.weightedStake,
+            });
+        });
+
+        it('reverts if the prior epoch is unfinalized', async () => {
+            await addActivePoolAsync();
+            await testContract.endEpoch.awaitTransactionSuccessAsync();
+            const tx = testContract.endEpoch.awaitTransactionSuccessAsync();
+            const expectedError = new StakingRevertErrors.PreviousEpochNotFinalizedError(
+                INITIAL_EPOCH,
+                1,
+            );
+            return expect(tx).to.revertWith(expectedError);
+        });
+    });
+
+    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]);
+            const rewardsPaidEvents = getRewardsPaidEvents(receipt.logs);
+            expect(rewardsPaidEvents.length).to.eq(1);
+            expect(rewardsPaidEvents[0].epoch).to.bignumber.eq(INITIAL_EPOCH + 1);
+            expect(rewardsPaidEvents[0].poolId).to.eq(pool.poolId);
+            assertEpochFinalizedEvent(
+                receipt.logs,
+                {
+                    epoch: new BigNumber(INITIAL_EPOCH),
+                    rewardsPaid: INITIAL_BALANCE,
+                },
+            );
+            assertDepositIntoStakingPoolRewardVaultCallEvent(
+                receipt.logs,
+                INITIAL_BALANCE,
+            );
+        });
+
+        it('can finalize multiple pools', async () => {
+            const nextEpoch = INITIAL_EPOCH + 1;
+            const pools = await Promise.all(_.times(3, () => addActivePoolAsync()));
+            const poolIds = pools.map(p => p.poolId);
+            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 [pool, event] of _.zip(pools, rewardsPaidEvents) as
+                    Array<[ActivePoolOpts, IStakingEventsRewardsPaidEventArgs]>) {
+                expect(event.epoch).to.bignumber.eq(nextEpoch);
+                expect(event.poolId).to.eq(pool.poolId);
+            }
+            assertEpochFinalizedEvent(
+                receipt.logs,
+                { epoch: new BigNumber(INITIAL_EPOCH) },
+            );
+            assertDepositIntoStakingPoolRewardVaultCallEvent(receipt.logs);
+        });
+
+        it('ignores a non-active pool', async () => {
+            const pools = await Promise.all(_.times(3, () => 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, () => 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, () => addActivePoolAsync()));
+            const pool = _.sample(pools) as ActivePoolOpts;
+            await testContract.endEpoch.awaitTransactionSuccessAsync();
+            await testContract.finalizePools.awaitTransactionSuccessAsync([pool.poolId]);
+            const poolState = await testContract
+                .internalGetActivePoolFromEpoch
+                .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);
+        });
+    });
+
+    describe('lifecycle', () => {
+        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 testContract.endEpoch.awaitTransactionSuccessAsync();
+            return expect(getCurrentEpochAsync()).to.become(INITIAL_EPOCH + 2);
+        });
+
+        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 addActivePoolAsync();
+            await testContract.endEpoch.awaitTransactionSuccessAsync();
+            expect(getCurrentEpochAsync()).to.become(INITIAL_EPOCH + 2);
+            const receipt = await testContract.finalizePools.awaitTransactionSuccessAsync([pool1.poolId]);
+            const rewardsPaidEvents = getRewardsPaidEvents(receipt.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 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 rewardsPaidEvents = getRewardsPaidEvents(receipt.logs);
+            expect(rewardsPaidEvents).to.deep.eq([]);
+        });
+    });
+
+    interface PoolRewards {
+        operatorReward: Numberish;
+        membersReward: Numberish;
+        membersStake: Numberish;
+    }
+
+    function assertPoolRewards(actual: PoolRewards, expected: Partial<PoolRewards>): void {
+        if (expected.operatorReward !== undefined) {
+            expect(actual.operatorReward).to.bignumber.eq(actual.operatorReward);
+        }
+        if (expected.membersReward !== undefined) {
+            expect(actual.membersReward).to.bignumber.eq(actual.membersReward);
+        }
+        if (expected.membersStake !== undefined) {
+            expect(actual.membersStake).to.bignumber.eq(actual.membersStake);
+        }
+    }
+
+    describe('_getUnfinalizedPoolReward()', () => {
+        const ZERO_REWARDS = {
+            operatorReward: 0,
+            membersReward: 0,
+            membersStake: 0,
+        };
+
+        it('returns empty if epoch is 0', async () => {
+            const poolId = hexRandom();
+            const rewards = await testContract
+                .internalGetUnfinalizedPoolRewards.callAsync(poolId);
+            assertPoolRewards(rewards, ZERO_REWARDS);
+        });
+
+        it('returns empty if pool was not active', async () => {
+            await testContract.endEpoch.awaitTransactionSuccessAsync();
+            const poolId = hexRandom();
+            const rewards = await testContract
+                .internalGetUnfinalizedPoolRewards.callAsync(poolId);
+            assertPoolRewards(rewards, ZERO_REWARDS);
+        });
+
+        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]);
+            const rewards = await testContract
+                .internalGetUnfinalizedPoolRewards.callAsync(pool.poolId);
+            assertPoolRewards(rewards, ZERO_REWARDS);
+        });
+
+        it('returns empty if pool was already finalized', async () => {
+            const pools = await Promise.all(_.times(3, () => addActivePoolAsync()));
+            const pool = _.sample(pools) as ActivePoolOpts;
+            await testContract.endEpoch.awaitTransactionSuccessAsync();
+            await testContract.finalizePools.awaitTransactionSuccessAsync([pool.poolId]);
+            const rewards = await testContract
+                .internalGetUnfinalizedPoolRewards.callAsync(pool.poolId);
+            assertPoolRewards(rewards, ZERO_REWARDS);
         });
     });
 });