From c4ca72cf224f6337dc1b25cdee94b0c7259d26a8 Mon Sep 17 00:00:00 2001
From: Greg Hysen <greg.hysen@gmail.com>
Date: Wed, 22 May 2019 14:17:23 -0700
Subject: [PATCH] staking end-to-end works

---
 .../staking/contracts/src/vaults/ZrxVault.sol |  4 +--
 contracts/staking/test/core_test.ts           | 31 +++++++++++--------
 .../staking/test/utils/staking_wrapper.ts     |  4 +++
 3 files changed, 23 insertions(+), 16 deletions(-)

diff --git a/contracts/staking/contracts/src/vaults/ZrxVault.sol b/contracts/staking/contracts/src/vaults/ZrxVault.sol
index 46d9c6cd68..660387ccf9 100644
--- a/contracts/staking/contracts/src/vaults/ZrxVault.sol
+++ b/contracts/staking/contracts/src/vaults/ZrxVault.sol
@@ -62,15 +62,13 @@ contract ZrxVault is
         external
         onlyStakingContract
     {
-        /*
         // deposit ZRX from owner
         erc20Proxy.transferFrom(
             zrxAssetData,
-            address(this),
             owner,
+            address(this),
             amount
         );
-        */
 
         // update balance
         balances[owner] = _safeAdd(balances[owner], amount);
diff --git a/contracts/staking/test/core_test.ts b/contracts/staking/test/core_test.ts
index 6f4128ec5a..21ab3b3954 100644
--- a/contracts/staking/test/core_test.ts
+++ b/contracts/staking/test/core_test.ts
@@ -71,22 +71,27 @@ describe('Staking Core', () => {
     });
     describe('end-to-end tests', () => {
         it('staking', async () => {
+            // setup tets parameters
             const amount = stakingWrapper.toBaseUnitAmount(10);
-            // check zrx token balance
-            const zrxTokenBalance = await stakingWrapper.getZrxTokenBalance(stakers[0]);
-            console.log(zrxTokenBalance);
-            expect(zrxTokenBalance).to.be.bignumber.gte(amount);
+            // check zrx token balances before minting stake
+            const zrxTokenBalanceOfVaultBeforeStaking = await stakingWrapper.getZrxTokenBalanceOfZrxVault();
+            expect(zrxTokenBalanceOfVaultBeforeStaking).to.be.bignumber.equal(new BigNumber(0));
+            const zrxTokenBalanceOfStakerBeforeStaking = await stakingWrapper.getZrxTokenBalance(stakers[0]);
+            expect(zrxTokenBalanceOfStakerBeforeStaking).to.be.bignumber.gte(amount);
             // mint stake
             const stakeMinted = await stakingWrapper.stake(stakers[0], amount);
-            console.log(stakeMinted);
-            //expect(stakeMinted).to.be.equal(amount);
-            
-            /*
-            const stakeBalance = stakingWrapper.getStakeBalance(stakers[0]);
-            expect(stakeBalance).to.be.equal(stakeAmount);
-            const vaultBalance = stakingWrapper.getZrxVaultBalance(stakers[0]);
-            expect(vaultBalance).to.be.equal(stakeAmount);
-            */
+            expect(stakeMinted).to.be.bignumber.equal(amount);
+            // check stake balance after minting
+            const stakeBalance = await stakingWrapper.getStakeBalance(stakers[0]);
+            expect(stakeBalance).to.be.bignumber.equal(amount);
+            // check zrx vault balance
+            const vaultBalance = await stakingWrapper.getZrxVaultBalance(stakers[0]);
+            expect(vaultBalance).to.be.bignumber.equal(amount);
+            // check zrx token balances
+            const zrxTokenBalanceOfVaultAfterStaking = await stakingWrapper.getZrxTokenBalanceOfZrxVault();
+            expect(zrxTokenBalanceOfVaultAfterStaking).to.be.bignumber.equal(amount);
+            const zrxTokenBalanceOfStakerAfterStaking = await stakingWrapper.getZrxTokenBalance(stakers[0]);
+            expect(zrxTokenBalanceOfStakerAfterStaking).to.be.bignumber.equal(zrxTokenBalanceOfStakerBeforeStaking.minus(amount));
         });
     });
 });
diff --git a/contracts/staking/test/utils/staking_wrapper.ts b/contracts/staking/test/utils/staking_wrapper.ts
index 5e4c9429b2..e4463da366 100644
--- a/contracts/staking/test/utils/staking_wrapper.ts
+++ b/contracts/staking/test/utils/staking_wrapper.ts
@@ -77,6 +77,10 @@ export class StakingWrapper {
         const balance = await this._zrxTokenContract.balanceOf.callAsync(holder);
         return balance;
     }
+    public async getZrxTokenBalanceOfZrxVault(): Promise<BigNumber> {
+        const balance = await this._zrxTokenContract.balanceOf.callAsync(this.getZrxVaultContract().address);
+        return balance;
+    }
     public toBaseUnitAmount(amount: BigNumber | number): BigNumber {
         const decimals = 18;
         const amountAsBigNumber = typeof(amount)  === 'number' ? new BigNumber(amount) : amount;