Add assertion generators to keeper, staker, taker mixins for the new function assertions

contracts/integrations/test/framework/actors/keeper.ts

@@ -1,12 +1,17 @@
 import {
+    AggregatedStats,
     IStakingEventsStakingPoolEarnedRewardsInEpochEventArgs,
-    TestStakingContract,
     TestStakingEvents,
 } from '@0x/contracts-staking';
 import { filterLogsToArguments, web3Wrapper } from '@0x/contracts-test-utils';
 import { BigNumber } from '@0x/utils';
 import { BlockParamLiteral, TransactionReceiptWithDecodedLogs } from 'ethereum-types';
 
+import { validEndEpochAssertion } from '../assertions/endEpoch';
+import { validFinalizePoolAssertion } from '../assertions/finalizePool';
+import { AssertionResult } from '../assertions/function_assertion';
+import { Pseudorandom } from '../utils/pseudorandom';
+
 import { Actor, Constructor } from './base';
 
 export interface KeeperInterface {
@@ -14,17 +19,6 @@ export interface KeeperInterface {
     finalizePoolsAsync: (poolIds?: string[]) => Promise<TransactionReceiptWithDecodedLogs[]>;
 }
 
-async function fastForwardToNextEpochAsync(stakingContract: TestStakingContract): Promise<void> {
-    // increase timestamp of next block by how many seconds we need to
-    // get to the next epoch.
-    const epochEndTime = await stakingContract.getCurrentEpochEarliestEndTimeInSeconds().callAsync();
-    const lastBlockTime = await web3Wrapper.getBlockTimestampAsync('latest');
-    const dt = Math.max(0, epochEndTime.minus(lastBlockTime).toNumber());
-    await web3Wrapper.increaseTimeAsync(dt);
-    // mine next block
-    await web3Wrapper.mineBlockAsync();
-}
-
 /**
  * This mixin encapsulates functionality associated with keepers within the 0x ecosystem.
  * This includes ending epochs sand finalizing pools in the staking system.
@@ -42,6 +36,13 @@ export function KeeperMixin<TBase extends Constructor>(Base: TBase): TBase & Con
             // tslint:disable-next-line:no-inferred-empty-object-type
             super(...args);
             this.actor = (this as any) as Actor;
+
+            // Register this mixin's assertion generators
+            this.actor.simulationActions = {
+                ...this.actor.simulationActions,
+                validFinalizePool: this._validFinalizePool(),
+                validEndEpoch: this._validEndEpoch(),
+            };
         }
 
         /**
@@ -50,7 +51,7 @@ export function KeeperMixin<TBase extends Constructor>(Base: TBase): TBase & Con
         public async endEpochAsync(shouldFastForward: boolean = true): Promise<TransactionReceiptWithDecodedLogs> {
             const { stakingWrapper } = this.actor.deployment.staking;
             if (shouldFastForward) {
-                await fastForwardToNextEpochAsync(stakingWrapper);
+                await this._fastForwardToNextEpochAsync();
             }
             return stakingWrapper.endEpoch().awaitTransactionSuccessAsync({ from: this.actor.address });
         }
@@ -83,6 +84,50 @@ export function KeeperMixin<TBase extends Constructor>(Base: TBase): TBase & Con
                 ),
             );
         }
+
+        private async *_validFinalizePool(): AsyncIterableIterator<AssertionResult | void> {
+            const { stakingPools } = this.actor.simulationEnvironment!;
+            const assertion = validFinalizePoolAssertion(this.actor.deployment, this.actor.simulationEnvironment!);
+            while (true) {
+                const poolId = Pseudorandom.sample(Object.keys(stakingPools));
+                if (poolId === undefined) {
+                    yield;
+                } else {
+                    yield assertion.executeAsync([poolId], { from: this.actor.address });
+                }
+            }
+        }
+
+        private async *_validEndEpoch(): AsyncIterableIterator<AssertionResult | void> {
+            const assertion = validEndEpochAssertion(this.actor.deployment, this.actor.simulationEnvironment!);
+            const { currentEpoch } = this.actor.simulationEnvironment!;
+            const { stakingWrapper } = this.actor.deployment.staking;
+            while (true) {
+                const aggregatedStats = AggregatedStats.fromArray(
+                    await stakingWrapper.aggregatedStatsByEpoch(currentEpoch.minus(1)).callAsync(),
+                );
+                if (aggregatedStats.numPoolsToFinalize.isGreaterThan(0)) {
+                    // Can't end the epoch if the previous epoch is not fully finalized.
+                    yield;
+                } else {
+                    await this._fastForwardToNextEpochAsync();
+                    yield assertion.executeAsync([], { from: this.actor.address });
+                }
+            }
+        }
+
+        private async _fastForwardToNextEpochAsync(): Promise<void> {
+            const { stakingWrapper } = this.actor.deployment.staking;
+
+            // increase timestamp of next block by how many seconds we need to
+            // get to the next epoch.
+            const epochEndTime = await stakingWrapper.getCurrentEpochEarliestEndTimeInSeconds().callAsync();
+            const lastBlockTime = await web3Wrapper.getBlockTimestampAsync('latest');
+            const dt = Math.max(0, epochEndTime.minus(lastBlockTime).toNumber());
+            await web3Wrapper.increaseTimeAsync(dt);
+            // mine next block
+            await web3Wrapper.mineBlockAsync();
+        }
     };
 }
 

contracts/integrations/test/framework/actors/staker.ts

@@ -7,6 +7,7 @@ import { AssertionResult } from '../assertions/function_assertion';
 import { validMoveStakeAssertion } from '../assertions/moveStake';
 import { validStakeAssertion } from '../assertions/stake';
 import { validUnstakeAssertion } from '../assertions/unstake';
+import { validWithdrawDelegatorRewardsAssertion } from '../assertions/withdrawDelegatorRewards';
 import { Pseudorandom } from '../utils/pseudorandom';
 
 import { Actor, Constructor } from './base';
@@ -44,6 +45,7 @@ export function StakerMixin<TBase extends Constructor>(Base: TBase): TBase & Con
                 validStake: this._validStake(),
                 validUnstake: this._validUnstake(),
                 validMoveStake: this._validMoveStake(),
+                validWithdrawDelegatorRewards: this._validWithdrawDelegatorRewards(),
             };
         }
 
@@ -129,6 +131,19 @@ export function StakerMixin<TBase extends Constructor>(Base: TBase): TBase & Con
                 yield assertion.executeAsync([from, to, amount], { from: this.actor.address });
             }
         }
+
+        private async *_validWithdrawDelegatorRewards(): AsyncIterableIterator<AssertionResult | void> {
+            const { stakingPools } = this.actor.simulationEnvironment!;
+            const assertion = validWithdrawDelegatorRewardsAssertion(this.actor.deployment, this.actor.simulationEnvironment!);
+            while (true) {
+                const poolId = Pseudorandom.sample(Object.keys(stakingPools));
+                if (poolId === undefined) {
+                    yield;
+                } else {
+                    yield assertion.executeAsync([poolId], { from: this.actor.address });
+                }
+            }
+        }
     };
 }
 

contracts/integrations/test/framework/actors/taker.ts

@@ -126,6 +126,7 @@ export function TakerMixin<TBase extends Constructor>(Base: TBase): TBase & Cons
                     yield assertion.executeAsync([order, fillAmount, order.signature], {
                         from: this.actor.address,
                     });
+                    // TODO: Randomly choose msg.value
                 }
             }
         }