From fd35249de868f3136c5c5c49ead072621d0ce20c Mon Sep 17 00:00:00 2001
From: Amir Bandeali <abandeali1@gmail.com>
Date: Mon, 23 Sep 2019 14:21:13 -0700
Subject: [PATCH] Track WETH reserved for rewards

---
 contracts/staking/contracts/src/immutable/MixinStorage.sol    | 3 +++
 contracts/staking/contracts/src/stake/MixinStake.sol          | 2 ++
 .../contracts/src/staking_pools/MixinStakingPoolRewards.sol   | 2 ++
 contracts/staking/contracts/src/sys/MixinFinalizer.sol        | 4 +++-
 contracts/staking/test/rewards_test.ts                        | 2 +-
 5 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/contracts/staking/contracts/src/immutable/MixinStorage.sol b/contracts/staking/contracts/src/immutable/MixinStorage.sol
index 7778e35483..ac3a268762 100644
--- a/contracts/staking/contracts/src/immutable/MixinStorage.sol
+++ b/contracts/staking/contracts/src/immutable/MixinStorage.sol
@@ -146,6 +146,9 @@ contract MixinStorage is
     /// @dev State for unfinalized rewards.
     IStructs.UnfinalizedState public unfinalizedState;
 
+    /// @dev The WETH balance of this contract that is reserved for pool reward payouts.
+    uint256 _reservedWethBalance;
+
     /// @dev Adds owner as an authorized address.
     constructor()
         public
diff --git a/contracts/staking/contracts/src/stake/MixinStake.sol b/contracts/staking/contracts/src/stake/MixinStake.sol
index 38c7671f61..7e9d727510 100644
--- a/contracts/staking/contracts/src/stake/MixinStake.sol
+++ b/contracts/staking/contracts/src/stake/MixinStake.sol
@@ -217,6 +217,7 @@ contract MixinStake is
         // to.
         IStructs.StoredBalance memory finalDelegatedStakeToPoolByOwner =
             _loadSyncedBalance(_delegatedStakeToPoolByOwner[owner][poolId]);
+
         _withdrawAndSyncDelegatorRewards(
             poolId,
             owner,
@@ -256,6 +257,7 @@ contract MixinStake is
         // from
         IStructs.StoredBalance memory finalDelegatedStakeToPoolByOwner =
             _loadSyncedBalance(_delegatedStakeToPoolByOwner[owner][poolId]);
+
         _withdrawAndSyncDelegatorRewards(
             poolId,
             owner,
diff --git a/contracts/staking/contracts/src/staking_pools/MixinStakingPoolRewards.sol b/contracts/staking/contracts/src/staking_pools/MixinStakingPoolRewards.sol
index 5837954811..b2c75d5c77 100644
--- a/contracts/staking/contracts/src/staking_pools/MixinStakingPoolRewards.sol
+++ b/contracts/staking/contracts/src/staking_pools/MixinStakingPoolRewards.sol
@@ -204,6 +204,7 @@ contract MixinStakingPoolRewards is
         if (membersReward > 0) {
             // Increment the balance of the pool
             balanceByPoolId[poolId] = balanceByPoolId[poolId].safeAdd(membersReward);
+            _reservedWethBalance = _reservedWethBalance.safeAdd(membersReward);
 
             // Fetch the last epoch at which we stored an entry for this pool;
             // this is the most up-to-date cumulative rewards for this pool.
@@ -297,6 +298,7 @@ contract MixinStakingPoolRewards is
 
         // Decrement the balance of the pool
         balanceByPoolId[poolId] = balanceByPoolId[poolId].safeSub(balance);
+        _reservedWethBalance = _reservedWethBalance.safeSub(balance);
 
         // Withdraw the member's WETH balance
         _getWethContract().transfer(member, balance);
diff --git a/contracts/staking/contracts/src/sys/MixinFinalizer.sol b/contracts/staking/contracts/src/sys/MixinFinalizer.sol
index 5ce7eb7bbc..630b84023c 100644
--- a/contracts/staking/contracts/src/sys/MixinFinalizer.sol
+++ b/contracts/staking/contracts/src/sys/MixinFinalizer.sol
@@ -248,7 +248,9 @@ contract MixinFinalizer is
         if (ethBalance != 0) {
             wethContract.deposit.value(ethBalance)();
         }
-        wethBalance = wethContract.balanceOf(address(this));
+        wethBalance = wethContract.balanceOf(address(this))
+            .safeSub(_reservedWethBalance);
+
         return wethBalance;
     }
 
diff --git a/contracts/staking/test/rewards_test.ts b/contracts/staking/test/rewards_test.ts
index 5179fddafa..19ada9994b 100644
--- a/contracts/staking/test/rewards_test.ts
+++ b/contracts/staking/test/rewards_test.ts
@@ -12,7 +12,7 @@ import { DelegatorsByPoolId, OperatorByPoolId, StakeInfo, StakeStatus } from './
 
 // tslint:disable:no-unnecessary-type-assertion
 // tslint:disable:max-file-line-count
-blockchainTests.resets.skip('Testing Rewards', env => {
+blockchainTests.resets('Testing Rewards', env => {
     // tokens & addresses
     let accounts: string[];
     let owner: string;