From 544e09cf4b152bac44a9eb8372d3024ed6b50446 Mon Sep 17 00:00:00 2001
From: Michael Zhu <mchl.zhu.96@gmail.com>
Date: Fri, 3 Jan 2020 15:46:58 -0500
Subject: [PATCH] invalidUnstakeAssertion

---
 .../test/framework/actors/staker.ts           | 22 ++++++++++++++-
 .../test/framework/assertions/unstake.ts      | 27 ++++++++++++++++++-
 2 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/contracts/integrations/test/framework/actors/staker.ts b/contracts/integrations/test/framework/actors/staker.ts
index 1f80abd913..425c359b35 100644
--- a/contracts/integrations/test/framework/actors/staker.ts
+++ b/contracts/integrations/test/framework/actors/staker.ts
@@ -8,7 +8,7 @@ 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';
+import { invalidUnstakeAssertion, validUnstakeAssertion } from '../assertions/unstake';
 import { validWithdrawDelegatorRewardsAssertion } from '../assertions/withdrawDelegatorRewards';
 import { Pseudorandom } from '../utils/pseudorandom';
 
@@ -49,6 +49,7 @@ export function StakerMixin<TBase extends Constructor>(Base: TBase): TBase & Con
                 validStake: this._validStake(),
                 invalidStake: this._invalidStake(),
                 validUnstake: this._validUnstake(),
+                invalidUnstake: this._invalidUnstake(),
                 validMoveStake: this._validMoveStake(),
                 validWithdrawDelegatorRewards: this._validWithdrawDelegatorRewards(),
             };
@@ -119,6 +120,25 @@ export function StakerMixin<TBase extends Constructor>(Base: TBase): TBase & Con
             }
         }
 
+        private async *_invalidUnstake(): AsyncIterableIterator<AssertionResult> {
+            const { stakingWrapper } = this.actor.deployment.staking;
+            const { deployment, balanceStore } = this.actor.simulationEnvironment!;
+
+            while (true) {
+                await balanceStore.updateErc20BalancesAsync();
+                const undelegatedStake = await stakingWrapper
+                    .getOwnerStakeByStatus(this.actor.address, StakeStatus.Undelegated)
+                    .callAsync();
+                const withdrawableStake = BigNumber.min(
+                    undelegatedStake.currentEpochBalance,
+                    undelegatedStake.nextEpochBalance,
+                );
+                const assertion = invalidUnstakeAssertion(deployment, withdrawableStake);
+                const amount = Pseudorandom.integer(withdrawableStake, constants.MAX_UINT256);
+                yield assertion.executeAsync([amount], { from: this.actor.address });
+            }
+        }
+
         private async *_validMoveStake(): AsyncIterableIterator<AssertionResult> {
             const { deployment, stakingPools } = this.actor.simulationEnvironment!;
             const assertion = validMoveStakeAssertion(deployment, this.actor.simulationEnvironment!, this.stake);
diff --git a/contracts/integrations/test/framework/assertions/unstake.ts b/contracts/integrations/test/framework/assertions/unstake.ts
index 0651da75f2..c140481bc5 100644
--- a/contracts/integrations/test/framework/assertions/unstake.ts
+++ b/contracts/integrations/test/framework/assertions/unstake.ts
@@ -1,4 +1,9 @@
-import { decreaseCurrentAndNextBalance, OwnerStakeByStatus, StakeStatus } from '@0x/contracts-staking';
+import {
+    decreaseCurrentAndNextBalance,
+    OwnerStakeByStatus,
+    StakeStatus,
+    StakingRevertErrors,
+} from '@0x/contracts-staking';
 import { expect } from '@0x/contracts-test-utils';
 import { BigNumber } from '@0x/utils';
 import { TxData } from 'ethereum-types';
@@ -81,4 +86,24 @@ export function validUnstakeAssertion(
         },
     });
 }
+
+/**
+ * Returns a FunctionAssertion for `unstake` which assumes that the input exceeds the amount that
+ * can be unstaked. Checks that the call reverts with an InsufficientBalanceError. Note that we
+ * close over `withdrawableStake` to avoid duplicating work done in the assertion generator.
+ */
+export function invalidUnstakeAssertion(
+    deployment: DeploymentManager,
+    withdrawableStake: BigNumber,
+): FunctionAssertion<[BigNumber], void, void> {
+    return new FunctionAssertion<[BigNumber], void, void>(deployment.staking.stakingWrapper, 'unstake', {
+        after: async (_beforeInfo: void, result: FunctionResult, args: [BigNumber]) => {
+            // Ensure that the tx reverted.
+            expect(result.success).to.be.false();
+
+            const [amount] = args;
+            expect(result.data).to.equal(new StakingRevertErrors.InsufficientBalanceError(amount, withdrawableStake));
+        },
+    });
+}
 /* tslint:enable:no-unnecessary-type-assertion */