Cleaned up staking contract with wrappers

contracts/staking/contracts/src/Staking.sol

@@ -18,440 +18,25 @@
 
 pragma solidity ^0.5.9;
 
-import "./immutable/MixinConstants.sol";
-import "./immutable/MixinStorage.sol";
-import "./core/MixinStake.sol";
-import "./core/MixinPools.sol";
-import "./core/MixinEpoch.sol";
-import "./core/MixinRewards.sol";
-import "./core/MixinFees.sol";
-import "./core/MixinExchange.sol";
+import "./wrappers/MixinStakeWrapper.sol";
+import "./wrappers/MixinStakeBalancesWrapper.sol";
+import "./wrappers/MixinPoolsWrapper.sol";
+import "./wrappers/MixinEpochWrapper.sol";
+import "./wrappers/MixinRewardsWrapper.sol";
+import "./wrappers/MixinFeesWrapper.sol";
+import "./wrappers/MixinExchangeWrapper.sol";
 
 
 contract Staking is
-    //IStaking,
-    //IStakingEvents,
-    MixinConstants,
-    MixinStorage,
-    MixinExchange,
-    MixinEpoch,
-    MixinRewards,
-    MixinStake,
-    MixinPools,
-    MixinFees
+    MixinExchangeWrapper,
+    MixinEpochWrapper,
+    MixinRewardsWrapper,
+    MixinStakeBalancesWrapper,
+    MixinStakeWrapper,
+    MixinPoolsWrapper,
+    MixinFeesWrapper
 {
 
-    ///// STAKE /////
-
-    function deposit(uint256 amount)
-        external
-    {
-        _deposit(msg.sender, amount);
-    }
-
-    function depositAndStake(uint256 amount)
-        external
-    {
-        _depositAndStake(msg.sender, amount);
-    }
-
-    function depositAndDelegate(bytes32 poolId, uint256 amount)
-        external
-    {
-        _depositAndDelegate(msg.sender, poolId, amount);
-    }
-
-    function activateStake(uint256 amount)
-        external
-    {
-        _activateStake(msg.sender, amount);
-    }
-
-    function activateAndDelegateStake(bytes32 poolId, uint256 amount)
-        external
-    {
-        _activateAndDelegateStake(msg.sender, poolId, amount);
-    }
-
-    function deactivateAndTimelockStake(uint256 amount)
-        external
-    {
-        _deactivateAndTimelockStake(msg.sender, amount);
-    }
-
-    function deactivateAndTimelockDelegatedStake(bytes32 poolId, uint256 amount)
-        external
-    {
-        _deactivateAndTimelockDelegatedStake(msg.sender, poolId, amount);
-    }
-
-    function withdraw(uint256 amount)
-        external
-    {
-        _withdraw(msg.sender, amount);
-    }
-
-    function forceTimelockSync(address owner)
-        external
-    {
-        _forceTimelockSync(owner);
-    }
-
-    ///// STAKE BALANCES /////
-
-    function getTotalStake(address owner)
-        external
-        view
-        returns (uint256)
-    {
-        return _getTotalStake(owner);
-    }
-
-    function getActivatedStake(address owner)
-        external
-        view
-        returns (uint256)
-    {
-        return _getActivatedStake(owner);
-    }
-
-    function getDeactivatedStake(address owner)
-        external
-        view
-        returns (uint256)
-    {
-        return _getDeactivatedStake(owner);
-    }
-
-    function getActivatableStake(address owner)
-        external
-        view
-        returns (uint256)
-    {
-        return _getActivatableStake(owner);
-    }
-
-    function getWithdrawableStake(address owner)
-        external
-        view
-        returns (uint256)
-    {
-        return _getWithdrawableStake(owner);
-    }
-
-    function getTimelockedStake(address owner)
-        external
-        view
-        returns (uint256)
-    {
-        return _getTimelockedStake(owner);
-    }
-
-    function getStakeDelegatedByOwner(address owner)
-        external
-        view
-        returns (uint256)
-    {
-        return _getStakeDelegatedByOwner(owner);
-    }
-
-    function getStakeDelegatedToPoolByOwner(address owner, bytes32 poolId)
-        external
-        view
-        returns (uint256)
-    {
-        return _getStakeDelegatedToPoolByOwner(owner, poolId);
-    }
-
-    function getStakeDelegatedToPool(bytes32 poolId)
-        external
-        view
-        returns (uint256)
-    {
-        return _getStakeDelegatedToPool(poolId);
-    }
-
-
-    ///// POOLS /////
-    modifier onlyPoolOperator(bytes32 poolId) {
-        require(
-            msg.sender == _getPoolOperator(poolId),
-            "ONLY_CALLABLE_BY_POOL_OPERATOR"
-        );
-
-        _;
-    }
-    
-    function getNextPoolId()
-        external
-        returns (bytes32 nextPoolId)
-    {
-        nextPoolId = _getNextPoolId();
-        return nextPoolId;
-    }
-    
-    function createPool(uint8 operatorShare)
-        external
-        returns (bytes32 poolId)
-    {
-        poolId = _createPool(msg.sender, operatorShare);
-        return poolId;
-    }
-
-    function addMakerToPool(
-        bytes32 poolId,
-        address makerAddress,
-        bytes calldata makerSignature
-    )
-        external
-        onlyPoolOperator(poolId)
-    {
-        _addMakerToPool(
-            poolId,
-            makerAddress,
-            makerSignature,
-            msg.sender
-        );
-    }
-
-    function removeMakerFromPool(bytes32 poolId, address makerAddress)
-        external
-        onlyPoolOperator(poolId)
-    {
-        _removeMakerFromPool(
-            poolId,
-            makerAddress,
-            msg.sender
-        );
-    }
-
-    function getMakerPoolId(address makerAddress)
-        external
-        view
-        returns (bytes32 makerId)
-    {
-        makerId = _getMakerPoolId(makerAddress);
-        return makerId;
-    }
-
-    function getMakerAddressesForPool(bytes32 makerId)
-        external
-        view
-        returns (address[] memory makerAddresses)
-    {
-        makerAddresses = _getMakerAddressesForPool(makerId);
-        return makerAddresses;
-    }
-
-    ///// EPOCHS /////
-
-    // @TODO - MixinAuthorizable
-    function goToNextEpoch()
-        external
-    {
-        _payRebates();
-        _goToNextEpoch();
-    }
-
-    function getEpochPeriodInSeconds()
-        external
-        view
-        returns (uint64)
-    {
-        return _getEpochPeriodInSeconds();
-    }
-
-    function getTimelockPeriodInEpochs()
-        external
-        view
-        returns (uint64)
-    {
-        return _getTimelockPeriodInEpochs();
-    }
-
-    function getCurrentEpochStartTimeInSeconds()
-        external
-        view
-        returns (uint64)
-    {
-        return _getCurrentEpochStartTimeInSeconds();
-    }
-
-    function getCurrentTimelockPeriodStartEpoch()
-        external
-        view
-        returns (uint64)
-    {
-        return _getCurrentTimelockPeriodStartEpoch();
-    }
-
-    function getCurrentEpochEndTimeInSeconds()
-        external
-        view
-        returns (uint64)
-    {
-        return _getCurrentEpochEndTimeInSeconds();
-    }
-
-    function getCurrentTimelockPeriodEndEpoch()
-        external
-        view
-        returns (uint64)
-    {
-        return _getCurrentTimelockPeriodEndEpoch();
-    }
-
-    function getCurrentEpoch()
-        external
-        view
-        returns (uint64)
-    {
-        return _getCurrentEpoch();
-    }
-
-    function getCurrentTimelockPeriod()
-        external
-        view
-        returns (uint64)
-    {
-        return _getCurrentTimelockPeriod();
-    }
-
-    ///// REWARDS /////
-
-    function getRewardBalance(bytes32 poolId)
-        external
-        view
-        returns (uint256)
-    {
-        return _getRewardBalance(poolId);
-    }
-
-    function getRewardBalanceOfOperator(bytes32 poolId)
-        external
-        view
-        returns (uint256)
-    {
-        return _getRewardBalanceOfOperator(poolId);
-    }
-
-    function getRewardBalanceOfPool(bytes32 poolId)
-        external
-        view
-        returns (uint256)
-    {
-        return _getRewardBalanceOfPool(poolId);
-    }
-
-     function computeRewardBalance(bytes32 poolId, address owner)
-        external
-        view
-        returns (uint256)
-    {
-        return _computeRewardBalance(poolId, owner);
-    }
-
-    function getShadowBalanceByPoolId(bytes32 poolId)
-        external
-        view
-        returns (uint256)
-    {
-        return _getShadowBalanceByPoolId(poolId);
-    }
-
-    function getShadowBalanceInPoolByOwner(address owner, bytes32 poolId)
-        external
-        view
-        returns (uint256)
-    {
-        return _getShadowBalanceInPoolByOwner(owner, poolId);
-    }
-
-    function withdrawOperatorReward(bytes32 poolId, uint256 amount)
-        external
-        onlyPoolOperator(poolId)
-    {
-        _withdrawOperatorReward(poolId, amount);
-    }
-
-    function withdrawReward(bytes32 poolId, uint256 amount)
-        external
-    {
-        _withdrawReward(poolId, msg.sender, amount);
-    }
-
-    function withdrawTotalOperatorReward(bytes32 poolId)
-        external
-        onlyPoolOperator(poolId)
-        returns (uint256)
-    {
-        return _withdrawTotalOperatorReward(poolId);
-    }
-
-    function withdrawTotalReward(bytes32 poolId)
-        external
-        returns (uint256)
-    {
-        return _withdrawTotalReward(poolId, msg.sender);
-    }
-
-    ///// FEES /////
-    modifier onlyExchange() {
-        require(
-            _isValidExchangeAddress(msg.sender),
-            "ONLY_CALLABLE_BY_EXCHANGE"
-        );
-        _;
-    }
-    function payProtocolFee(address makerAddress)
-        external
-        payable
-        onlyExchange
-    {
-        _payProtocolFee(makerAddress, msg.value);
-    }
-
-    function getProtocolFeesThisEpochByPool(bytes32 poolId)
-        external
-        view
-        returns (uint256)
-    {
-        return _getProtocolFeesThisEpochByPool(poolId);
-    }
-
-    function getTotalProtocolFeesThisEpoch()
-        external
-        view
-        returns (uint256)
-    {
-        return _getTotalProtocolFeesThisEpoch();
-    }
-
-    ///// EXCHANGES /////
-    // @TODO - Only by 0x multi-sig
-
-    function isValidExchangeAddress(address addr)
-        external
-        view
-        returns (bool)
-    {
-        return _isValidExchangeAddress(addr);
-    }
-
-    function addExchangeAddress(address addr)
-        external
-    {
-        _addExchangeAddress(addr);
-    }
-
-    function removeExchangeAddress(address addr)
-        external
-    {
-        _removeExchangeAddress(addr);
-    }
-
-    ///// SETTERS /////
-
     function setZrxVault(address _zrxVault)
         external
     {

contracts/staking/contracts/src/core/MixinFees.sol

@@ -163,4 +163,11 @@ contract MixinFees is
             totalRewardsRecordedInVault
         );
     }
+
+    function _finalizeFees()
+        internal
+    {
+        _payRebates();
+        _goToNextEpoch();
+    }
 }

contracts/staking/contracts/src/utils/modifiers.sol

@@ -0,0 +1,46 @@
+/*
+
+  Copyright 2018 ZeroEx Intl.
+
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+
+*/
+
+pragma solidity ^0.5.5;
+
+import "../core/MixinPools.sol";
+import "../core/MixinExchange.sol";
+
+
+contract Modifiers is
+    MixinPools,
+    MixinExchange
+{
+    
+    modifier onlyPoolOperator(bytes32 poolId) {
+        require(
+            msg.sender == _getPoolOperator(poolId),
+            "ONLY_CALLABLE_BY_POOL_OPERATOR"
+        );
+
+        _;
+    }
+
+    modifier onlyExchange() {
+        require(
+            _isValidExchangeAddress(msg.sender),
+            "ONLY_CALLABLE_BY_EXCHANGE"
+        );
+        _;
+    }
+}

contracts/staking/contracts/src/wrappers/MixinEpochWrapper.sol

@@ -0,0 +1,92 @@
+/*
+
+  Copyright 2018 ZeroEx Intl.
+
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+
+*/
+
+pragma solidity ^0.5.5;
+
+import "../core/MixinEpoch.sol";
+
+
+contract MixinEpochWrapper is
+    MixinEpoch
+{
+    // @TODO - MixinAuthorizable
+
+    function getEpochPeriodInSeconds()
+        external
+        view
+        returns (uint64)
+    {
+        return _getEpochPeriodInSeconds();
+    }
+
+    function getTimelockPeriodInEpochs()
+        external
+        view
+        returns (uint64)
+    {
+        return _getTimelockPeriodInEpochs();
+    }
+
+    function getCurrentEpochStartTimeInSeconds()
+        external
+        view
+        returns (uint64)
+    {
+        return _getCurrentEpochStartTimeInSeconds();
+    }
+
+    function getCurrentTimelockPeriodStartEpoch()
+        external
+        view
+        returns (uint64)
+    {
+        return _getCurrentTimelockPeriodStartEpoch();
+    }
+
+    function getCurrentEpochEndTimeInSeconds()
+        external
+        view
+        returns (uint64)
+    {
+        return _getCurrentEpochEndTimeInSeconds();
+    }
+
+    function getCurrentTimelockPeriodEndEpoch()
+        external
+        view
+        returns (uint64)
+    {
+        return _getCurrentTimelockPeriodEndEpoch();
+    }
+
+    function getCurrentEpoch()
+        external
+        view
+        returns (uint64)
+    {
+        return _getCurrentEpoch();
+    }
+
+    function getCurrentTimelockPeriod()
+        external
+        view
+        returns (uint64)
+    {
+        return _getCurrentTimelockPeriod();
+    }
+}

contracts/staking/contracts/src/wrappers/MixinExchangeWrapper.sol

@@ -0,0 +1,49 @@
+/*
+
+  Copyright 2018 ZeroEx Intl.
+
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+
+*/
+
+pragma solidity ^0.5.5;
+
+import "../core/MixinExchange.sol";
+
+
+contract MixinExchangeWrapper is
+    MixinExchange
+{
+
+    // @TODO - Only by 0x multi-sig
+
+    function isValidExchangeAddress(address addr)
+        external
+        view
+        returns (bool)
+    {
+        return _isValidExchangeAddress(addr);
+    }
+
+    function addExchangeAddress(address addr)
+        external
+    {
+        _addExchangeAddress(addr);
+    }
+
+    function removeExchangeAddress(address addr)
+        external
+    {
+        _removeExchangeAddress(addr);
+    }
+}

contracts/staking/contracts/src/wrappers/MixinFeesWrapper.sol

@@ -0,0 +1,59 @@
+/*
+
+  Copyright 2018 ZeroEx Intl.
+
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+
+*/
+
+pragma solidity ^0.5.5;
+
+import "../core/MixinFees.sol";
+import "../utils/Modifiers.sol";
+
+
+contract MixinFeesWrapper is
+    MixinFees,
+    Modifiers
+{
+
+    function payProtocolFee(address makerAddress)
+        external
+        payable
+        onlyExchange
+    {
+        _payProtocolFee(makerAddress, msg.value);
+    }
+
+    function getProtocolFeesThisEpochByPool(bytes32 poolId)
+        external
+        view
+        returns (uint256)
+    {
+        return _getProtocolFeesThisEpochByPool(poolId);
+    }
+
+    function getTotalProtocolFeesThisEpoch()
+        external
+        view
+        returns (uint256)
+    {
+        return _getTotalProtocolFeesThisEpoch();
+    }
+
+    function finalizeFees()
+        external
+    {
+        _finalizeFees();
+    }
+}

contracts/staking/contracts/src/wrappers/MixinPoolsWrapper.sol

@@ -0,0 +1,90 @@
+/*
+
+  Copyright 2018 ZeroEx Intl.
+
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+
+*/
+
+pragma solidity ^0.5.5;
+
+import "../core/MixinPools.sol";
+import "../utils/Modifiers.sol";
+
+
+contract MixinPoolsWrapper is
+    MixinPools,
+    Modifiers
+{
+
+    function getNextPoolId()
+        external
+        returns (bytes32 nextPoolId)
+    {
+        nextPoolId = _getNextPoolId();
+        return nextPoolId;
+    }
+    
+    function createPool(uint8 operatorShare)
+        external
+        returns (bytes32 poolId)
+    {
+        poolId = _createPool(msg.sender, operatorShare);
+        return poolId;
+    }
+
+    function addMakerToPool(
+        bytes32 poolId,
+        address makerAddress,
+        bytes calldata makerSignature
+    )
+        external
+        onlyPoolOperator(poolId)
+    {
+        _addMakerToPool(
+            poolId,
+            makerAddress,
+            makerSignature,
+            msg.sender
+        );
+    }
+
+    function removeMakerFromPool(bytes32 poolId, address makerAddress)
+        external
+        onlyPoolOperator(poolId)
+    {
+        _removeMakerFromPool(
+            poolId,
+            makerAddress,
+            msg.sender
+        );
+    }
+
+    function getMakerPoolId(address makerAddress)
+        external
+        view
+        returns (bytes32 makerId)
+    {
+        makerId = _getMakerPoolId(makerAddress);
+        return makerId;
+    }
+
+    function getMakerAddressesForPool(bytes32 makerId)
+        external
+        view
+        returns (address[] memory makerAddresses)
+    {
+        makerAddresses = _getMakerAddressesForPool(makerId);
+        return makerAddresses;
+    }
+}

contracts/staking/contracts/src/wrappers/MixinRewardsWrapper.sol

@@ -0,0 +1,105 @@
+/*
+
+  Copyright 2018 ZeroEx Intl.
+
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+
+*/
+
+pragma solidity ^0.5.5;
+
+import "../core/MixinRewards.sol";
+import "../utils/Modifiers.sol";
+
+
+contract MixinRewardsWrapper is
+    MixinRewards,
+    Modifiers
+{
+
+    function getRewardBalance(bytes32 poolId)
+        external
+        view
+        returns (uint256)
+    {
+        return _getRewardBalance(poolId);
+    }
+
+    function getRewardBalanceOfOperator(bytes32 poolId)
+        external
+        view
+        returns (uint256)
+    {
+        return _getRewardBalanceOfOperator(poolId);
+    }
+
+    function getRewardBalanceOfPool(bytes32 poolId)
+        external
+        view
+        returns (uint256)
+    {
+        return _getRewardBalanceOfPool(poolId);
+    }
+
+     function computeRewardBalance(bytes32 poolId, address owner)
+        external
+        view
+        returns (uint256)
+    {
+        return _computeRewardBalance(poolId, owner);
+    }
+
+    function getShadowBalanceByPoolId(bytes32 poolId)
+        external
+        view
+        returns (uint256)
+    {
+        return _getShadowBalanceByPoolId(poolId);
+    }
+
+    function getShadowBalanceInPoolByOwner(address owner, bytes32 poolId)
+        external
+        view
+        returns (uint256)
+    {
+        return _getShadowBalanceInPoolByOwner(owner, poolId);
+    }
+
+    function withdrawOperatorReward(bytes32 poolId, uint256 amount)
+        external
+        onlyPoolOperator(poolId)
+    {
+        _withdrawOperatorReward(poolId, amount);
+    }
+
+    function withdrawReward(bytes32 poolId, uint256 amount)
+        external
+    {
+        _withdrawReward(poolId, msg.sender, amount);
+    }
+
+    function withdrawTotalOperatorReward(bytes32 poolId)
+        external
+        onlyPoolOperator(poolId)
+        returns (uint256)
+    {
+        return _withdrawTotalOperatorReward(poolId);
+    }
+
+    function withdrawTotalReward(bytes32 poolId)
+        external
+        returns (uint256)
+    {
+        return _withdrawTotalReward(poolId, msg.sender);
+    }
+}

contracts/staking/contracts/src/wrappers/MixinStakeBalancesWrapper.sol

@@ -0,0 +1,99 @@
+/*
+
+  Copyright 2018 ZeroEx Intl.
+
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+
+*/
+
+pragma solidity ^0.5.5;
+
+import "../core/MixinStakeBalances.sol";
+
+
+contract MixinStakeBalancesWrapper is
+    MixinStakeBalances
+{
+
+    function getTotalStake(address owner)
+        external
+        view
+        returns (uint256)
+    {
+        return _getTotalStake(owner);
+    }
+
+    function getActivatedStake(address owner)
+        external
+        view
+        returns (uint256)
+    {
+        return _getActivatedStake(owner);
+    }
+
+    function getDeactivatedStake(address owner)
+        external
+        view
+        returns (uint256)
+    {
+        return _getDeactivatedStake(owner);
+    }
+
+    function getActivatableStake(address owner)
+        external
+        view
+        returns (uint256)
+    {
+        return _getActivatableStake(owner);
+    }
+
+    function getWithdrawableStake(address owner)
+        external
+        view
+        returns (uint256)
+    {
+        return _getWithdrawableStake(owner);
+    }
+
+    function getTimelockedStake(address owner)
+        external
+        view
+        returns (uint256)
+    {
+        return _getTimelockedStake(owner);
+    }
+
+    function getStakeDelegatedByOwner(address owner)
+        external
+        view
+        returns (uint256)
+    {
+        return _getStakeDelegatedByOwner(owner);
+    }
+
+    function getStakeDelegatedToPoolByOwner(address owner, bytes32 poolId)
+        external
+        view
+        returns (uint256)
+    {
+        return _getStakeDelegatedToPoolByOwner(owner, poolId);
+    }
+
+    function getStakeDelegatedToPool(bytes32 poolId)
+        external
+        view
+        returns (uint256)
+    {
+        return _getStakeDelegatedToPool(poolId);
+    }
+}

contracts/staking/contracts/src/wrappers/MixinStakeWrapper.sol

@@ -0,0 +1,81 @@
+/*
+
+  Copyright 2018 ZeroEx Intl.
+
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+
+*/
+
+pragma solidity ^0.5.5;
+
+import "../core/MixinStake.sol";
+
+
+contract MixinStakeWrapper is
+    MixinStake
+{
+
+    function deposit(uint256 amount)
+        external
+    {
+        _deposit(msg.sender, amount);
+    }
+
+    function depositAndStake(uint256 amount)
+        external
+    {
+        _depositAndStake(msg.sender, amount);
+    }
+
+    function depositAndDelegate(bytes32 poolId, uint256 amount)
+        external
+    {
+        _depositAndDelegate(msg.sender, poolId, amount);
+    }
+
+    function activateStake(uint256 amount)
+        external
+    {
+        _activateStake(msg.sender, amount);
+    }
+
+    function activateAndDelegateStake(bytes32 poolId, uint256 amount)
+        external
+    {
+        _activateAndDelegateStake(msg.sender, poolId, amount);
+    }
+
+    function deactivateAndTimelockStake(uint256 amount)
+        external
+    {
+        _deactivateAndTimelockStake(msg.sender, amount);
+    }
+
+    function deactivateAndTimelockDelegatedStake(bytes32 poolId, uint256 amount)
+        external
+    {
+        _deactivateAndTimelockDelegatedStake(msg.sender, poolId, amount);
+    }
+
+    function withdraw(uint256 amount)
+        external
+    {
+        _withdraw(msg.sender, amount);
+    }
+
+    function forceTimelockSync(address owner)
+        external
+    {
+        _forceTimelockSync(owner);
+    }
+}

contracts/staking/test/utils/staking_wrapper.ts

@@ -274,7 +274,7 @@ export class StakingWrapper {
     }
     ///// EPOCHS /////
     public async goToNextEpochAsync(): Promise<TransactionReceiptWithDecodedLogs> {
-        const calldata = this.getStakingContract().goToNextEpoch.getABIEncodedTransactionData();
+        const calldata = this.getStakingContract().finalizeFees.getABIEncodedTransactionData();
         const txReceipt = await this._executeTransactionAsync(calldata, undefined, new BigNumber(0), true);
         console.log(JSON.stringify(txReceipt, null , 4));
         //console.log((txReceipt.logs[0] as LogWithDecodedArgs<StakingEEventArgs>).args);