From c110dc9e6ab8108f4b7c0da18411a9f32ec64562 Mon Sep 17 00:00:00 2001
From: Michael Zhu <mchl.zhu.96@gmail.com>
Date: Fri, 3 Jan 2020 15:27:24 -0500
Subject: [PATCH] generic assertion for TRANSFER_FAILED reverts; _invalidStake
 generator

---
 .../test/framework/actors/staker.ts           | 16 +++++++++++++
 .../assertions/generic_assertions.ts          | 24 +++++++++++++++++++
 2 files changed, 40 insertions(+)
 create mode 100644 contracts/integrations/test/framework/assertions/generic_assertions.ts

diff --git a/contracts/integrations/test/framework/actors/staker.ts b/contracts/integrations/test/framework/actors/staker.ts
index 52f9083668..1f80abd913 100644
--- a/contracts/integrations/test/framework/actors/staker.ts
+++ b/contracts/integrations/test/framework/actors/staker.ts
@@ -1,9 +1,11 @@
 import { OwnerStakeByStatus, StakeInfo, StakeStatus, StoredBalance } from '@0x/contracts-staking';
+import { constants } from '@0x/contracts-test-utils';
 import { BigNumber } from '@0x/utils';
 import '@azure/core-asynciterator-polyfill';
 import * as _ from 'lodash';
 
 import { AssertionResult } from '../assertions/function_assertion';
+import { assetProxyTransferFailedAssertion } from '../assertions/generic_assertions';
 import { validMoveStakeAssertion } from '../assertions/moveStake';
 import { validStakeAssertion } from '../assertions/stake';
 import { validUnstakeAssertion } from '../assertions/unstake';
@@ -45,6 +47,7 @@ export function StakerMixin<TBase extends Constructor>(Base: TBase): TBase & Con
             this.actor.simulationActions = {
                 ...this.actor.simulationActions,
                 validStake: this._validStake(),
+                invalidStake: this._invalidStake(),
                 validUnstake: this._validUnstake(),
                 validMoveStake: this._validMoveStake(),
                 validWithdrawDelegatorRewards: this._validWithdrawDelegatorRewards(),
@@ -84,6 +87,19 @@ export function StakerMixin<TBase extends Constructor>(Base: TBase): TBase & Con
             }
         }
 
+        private async *_invalidStake(): AsyncIterableIterator<AssertionResult> {
+            const { zrx } = this.actor.deployment.tokens;
+            const { deployment, balanceStore } = this.actor.simulationEnvironment!;
+            const assertion = assetProxyTransferFailedAssertion(deployment.staking.stakingWrapper, 'stake');
+
+            while (true) {
+                await balanceStore.updateErc20BalancesAsync();
+                const zrxBalance = balanceStore.balances.erc20[this.actor.address][zrx.address];
+                const amount = Pseudorandom.integer(zrxBalance, constants.MAX_UINT256);
+                yield assertion.executeAsync([amount], { from: this.actor.address });
+            }
+        }
+
         private async *_validUnstake(): AsyncIterableIterator<AssertionResult> {
             const { stakingWrapper } = this.actor.deployment.staking;
             const { deployment, balanceStore } = this.actor.simulationEnvironment!;
diff --git a/contracts/integrations/test/framework/assertions/generic_assertions.ts b/contracts/integrations/test/framework/assertions/generic_assertions.ts
new file mode 100644
index 0000000000..49febd2ac3
--- /dev/null
+++ b/contracts/integrations/test/framework/assertions/generic_assertions.ts
@@ -0,0 +1,24 @@
+import { BaseContract } from '@0x/base-contract';
+import { expect } from '@0x/contracts-test-utils';
+import { RevertReason } from '@0x/types';
+import { StringRevertError } from '@0x/utils';
+
+import { FunctionAssertion, FunctionResult } from './function_assertion';
+
+/**
+ * Returns a generic FunctionAssertion for the given contract function, asserting that the
+ * function call reverts in an asset proxy contract with TRANSFER_FAILED.
+ */
+export function assetProxyTransferFailedAssertion<TArgs extends any[]>(
+    contract: BaseContract,
+    functionName: string,
+): FunctionAssertion<TArgs, void, void> {
+    return new FunctionAssertion(contract, functionName, {
+        after: async (_beforeInfo: void, result: FunctionResult) => {
+            // Ensure that the tx reverted.
+            expect(result.success).to.be.false();
+            // Check revert error
+            expect(result.data).to.equal(new StringRevertError(RevertReason.TransferFailed));
+        },
+    });
+}