fixing bugs

contracts/integrations/test/framework/assertions/endEpoch.ts

@@ -1,8 +1,14 @@
 import { WETH9DepositEventArgs, WETH9Events } from '@0x/contracts-erc20';
-import { AggregatedStats, StakingEvents, StakingEpochEndedEventArgs } from '@0x/contracts-staking';
-import { expect, verifyEventsFromLogs } from '@0x/contracts-test-utils';
+import {
+    AggregatedStats,
+    StakingEvents,
+    StakingEpochEndedEventArgs,
+    StakingEpochFinalizedEventArgs,
+} from '@0x/contracts-staking';
+import { constants, expect, verifyEventsFromLogs } from '@0x/contracts-test-utils';
 import { BigNumber } from '@0x/utils';
 import { TxData } from 'ethereum-types';
+import * as _ from 'lodash';
 
 import { DeploymentManager } from '../deployment_manager';
 import { SimulationEnvironment } from '../simulation';
@@ -25,20 +31,25 @@ export function validEndEpochAssertion(
     simulationEnvironment: SimulationEnvironment,
 ): FunctionAssertion<[], EndEpochBeforeInfo, void> {
     const { stakingWrapper } = deployment.staking;
-    const { balanceStore, currentEpoch } = simulationEnvironment;
+    const { balanceStore } = simulationEnvironment;
 
     return new FunctionAssertion(stakingWrapper, 'endEpoch', {
         before: async () => {
             await balanceStore.updateEthBalancesAsync();
             const aggregatedStatsBefore = AggregatedStats.fromArray(
-                await stakingWrapper.aggregatedStatsByEpoch(currentEpoch).callAsync(),
+                await stakingWrapper.aggregatedStatsByEpoch(simulationEnvironment.currentEpoch).callAsync(),
             );
             const wethReservedForPoolRewards = await stakingWrapper.wethReservedForPoolRewards().callAsync();
             return { wethReservedForPoolRewards, aggregatedStatsBefore };
         },
         after: async (beforeInfo: EndEpochBeforeInfo, result: FunctionResult, _args: [], _txData: Partial<TxData>) => {
+            // Ensure that the tx succeeded.
+            expect(result.success, `Error: ${result.data}`).to.be.true();
+
+            const { currentEpoch } = simulationEnvironment;
+
             // Check WETH deposit event
-            const previousEthBalance = balanceStore.balances.eth[stakingWrapper.address];
+            const previousEthBalance = balanceStore.balances.eth[stakingWrapper.address] || constants.ZERO_AMOUNT;
             if (previousEthBalance.isGreaterThan(0)) {
                 verifyEventsFromLogs<WETH9DepositEventArgs>(
                     result.receipt!.logs,
@@ -56,9 +67,11 @@ export function validEndEpochAssertion(
             const { wethReservedForPoolRewards, aggregatedStatsBefore } = beforeInfo;
             const expectedAggregatedStats = {
                 ...aggregatedStatsBefore,
-                rewardsAvailable: balanceStore.balances.erc20[stakingWrapper.address][
-                    deployment.tokens.weth.address
-                ].minus(wethReservedForPoolRewards),
+                rewardsAvailable: _.get(
+                    balanceStore.balances,
+                    ['erc20', stakingWrapper.address, deployment.tokens.weth.address],
+                    constants.ZERO_AMOUNT,
+                ).minus(wethReservedForPoolRewards),
             };
 
             const aggregatedStatsAfter = AggregatedStats.fromArray(
@@ -66,22 +79,35 @@ export function validEndEpochAssertion(
             );
             expect(aggregatedStatsAfter).to.deep.equal(expectedAggregatedStats);
 
-            const expectedEpochEndedEvents = aggregatedStatsAfter.numPoolsToFinalize.isZero()
+            verifyEventsFromLogs<StakingEpochEndedEventArgs>(
+                result.receipt!.logs,
+                [
+                    {
+                        epoch: currentEpoch,
+                        numPoolsToFinalize: aggregatedStatsAfter.numPoolsToFinalize,
+                        rewardsAvailable: aggregatedStatsAfter.rewardsAvailable,
+                        totalFeesCollected: aggregatedStatsAfter.totalFeesCollected,
+                        totalWeightedStake: aggregatedStatsAfter.totalWeightedStake,
+                    },
+                ],
+                StakingEvents.EpochEnded,
+            );
+
+            const expectedEpochFinalizedEvents = aggregatedStatsAfter.numPoolsToFinalize.isZero()
                 ? [
                       {
                           epoch: currentEpoch,
-                          numPoolsToFinalize: aggregatedStatsAfter.numPoolsToFinalize,
-                          rewardsAvailable: aggregatedStatsAfter.rewardsAvailable,
-                          totalFeesCollected: aggregatedStatsAfter.totalFeesCollected,
-                          totalWeightedStake: aggregatedStatsAfter.totalWeightedStake,
+                          rewardsPaid: constants.ZERO_AMOUNT,
+                          rewardsRemaining: aggregatedStatsAfter.rewardsAvailable,
                       },
                   ]
                 : [];
-            verifyEventsFromLogs<StakingEpochEndedEventArgs>(
+            verifyEventsFromLogs<StakingEpochFinalizedEventArgs>(
                 result.receipt!.logs,
-                expectedEpochEndedEvents,
-                StakingEvents.EpochEnded,
+                expectedEpochFinalizedEvents,
+                StakingEvents.EpochFinalized,
             );
+
             expect(result.data, 'endEpoch should return the number of unfinalized pools').to.bignumber.equal(
                 aggregatedStatsAfter.numPoolsToFinalize,
             );