@0x:contracts-staking Added tests for protocol fees

contracts/staking/contracts/src/fees/MixinExchangeFees.sol

@@ -99,37 +99,42 @@ contract MixinExchangeFees is
         payable
         onlyExchange
     {
-        // Get the pool id of the maker address, and use this pool id to get the amount
+        // If the protocol fee payment is invalid, revert with a rich error.
-        // of fees collected during this epoch.
+        if (
-        bytes32 poolId = getStakingPoolIdOfMaker(makerAddress);
+            protocolFeePaid == 0 ||
-        uint256 _feesCollectedThisEpoch = protocolFeesThisEpochByPool[poolId];
+            (msg.value != protocolFeePaid && msg.value != 0)
-
+        ) {
-        if (msg.value == 0) {
-            // Transfer the protocol fee to this address.
-            erc20Proxy.transferFrom(
-                wethAssetData,
-                payerAddress,
-                address(this),
-                protocolFeePaid
-            );
-
-            // Update the amount of protocol fees paid to this pool this epoch.
-            protocolFeesThisEpochByPool[poolId] = _feesCollectedThisEpoch.safeAdd(protocolFeePaid);
-
-        } else if (msg.value == protocolFeePaid) {
-            // Update the amount of protocol fees paid to this pool this epoch.
-            protocolFeesThisEpochByPool[poolId] = _feesCollectedThisEpoch.safeAdd(protocolFeePaid);
-        } else {
-            // If the wrong message value was sent, revert with a rich error.
             LibRichErrors.rrevert(LibStakingRichErrors.InvalidProtocolFeePaymentError(
                 protocolFeePaid,
                 msg.value
             ));
         }
 
-        // If there were no fees collected prior to this payment, activate the pool that is being paid.
+        // Transfer the protocol fee to this address if it should be paid in WETH.
-        if (_feesCollectedThisEpoch == 0) {
+        if (msg.value == 0) {
-            activePoolsThisEpoch.push(poolId);
+            erc20Proxy.transferFrom(
+                WETH_ASSET_DATA,
+                payerAddress,
+                address(this),
+                protocolFeePaid
+            );
+        }
+
+        // Get the pool id of the maker address.
+        bytes32 poolId = getStakingPoolIdOfMaker(makerAddress);
+
+        // Only attribute the protocol fee payment to a pool if the maker is registered to a pool.
+        if (poolId != NIL_POOL_ID) {
+            // Use the maker pool id to get the amount of fees collected during this epoch in the pool.
+            uint256 _feesCollectedThisEpoch = protocolFeesThisEpochByPool[poolId];
+
+            // Update the amount of protocol fees paid to this pool this epoch.
+            protocolFeesThisEpochByPool[poolId] = _feesCollectedThisEpoch.safeAdd(protocolFeePaid);
+
+            // If there were no fees collected prior to this payment, activate the pool that is being paid.
+            if (_feesCollectedThisEpoch == 0) {
+                activePoolsThisEpoch.push(poolId);
+            }
         }
     }
 
@@ -179,6 +184,18 @@ contract MixinExchangeFees is
         return protocolFeesThisEpochByPool[poolId];
     }
 
+    /// @dev Withdraws the entire WETH balance of the contract.
+    function _unwrapWETH()
+        internal
+    {
+        uint256 wethBalance = IEtherToken(WETH_ADDRESS).balanceOf(address(this));
+
+        // Don't withdraw WETH if the WETH balance is zero as a gas optimization.
+        if (wethBalance != 0) {
+            IEtherToken(WETH_ADDRESS).withdraw(wethBalance);
+        }
+    }
+
     /// @dev Pays rewards to market making pools that were active this epoch.
     /// Each pool receives a portion of the fees generated this epoch (see _cobbDouglas) that is
     /// proportional to (i) the fee volume attributed to their pool over the epoch, and
@@ -204,11 +221,12 @@ contract MixinExchangeFees is
             uint256 finalContractBalance
         )
     {
-        // step 1/4 - withdraw the entire wrapper ether balance into this contract.
+        // step 1/4 - withdraw the entire wrapped ether balance into this contract. WETH
-        //            WETH is unwrapped here to keep `payProtocolFee()` calls relatively cheap.
+        //            is unwrapped here to keep `payProtocolFee()` calls relatively cheap,
-        uint256 wethBalance = IEtherToken(WETH_ADDRESS).balanceOf(address(this));
+        //            and WETH is only withdrawn if this contract's WETH balance is nonzero.
-        IEtherToken(WETH_ADDRESS).withdraw(wethBalance);
+        _unwrapWETH();
 
+        // Initialize initial values
         totalActivePools = activePoolsThisEpoch.length;
         totalFeesCollected = 0;
         totalWeightedStake = 0;

contracts/staking/contracts/src/immutable/MixinConstants.sol

@@ -34,6 +34,8 @@ contract MixinConstants is
 
     bytes32 constant internal NIL_POOL_ID = 0x0000000000000000000000000000000000000000000000000000000000000000;
 
+    bytes32 constant internal NIL_POOL_ID = 0x0000000000000000000000000000000000000000000000000000000000000000;
+
     address constant internal NIL_ADDRESS = 0x0000000000000000000000000000000000000000;
 
     bytes32 constant internal UNKNOWN_STAKING_POOL_ID = 0x0000000000000000000000000000000000000000000000000000000000000000;
@@ -46,4 +48,7 @@ contract MixinConstants is
 
     // The address of the canonical WETH contract -- this will be used as an alternative to ether for paying protocol fees.
     address constant internal WETH_ADDRESS = address(0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2);
+
+    // abi.encodeWithSelector(IAssetData(address(0)).ERC20Token.selector, WETH_ADDRESS)
+    bytes constant internal WETH_ASSET_DATA = hex"f47261b0000000000000000000000000c02aaa39b223fe8d0a0e5c4f27ead9083c756cc2";
 }

contracts/staking/contracts/src/immutable/MixinStorage.sol

@@ -37,20 +37,11 @@ contract MixinStorage is
     constructor()
         public
         Ownable()
-    {
+    {} // solhint-disable-line no-empty-blocks
-        // Set the erc20 asset proxy data.
-        wethAssetData = abi.encodeWithSelector(
-            IAssetData(address(0)).ERC20Token.selector,
-            WETH_ADDRESS
-        );
-    }
 
     // 0x ERC20 Proxy
     IAssetProxy internal erc20Proxy;
 
-    // The asset data that should be sent to transfer weth
-    bytes internal wethAssetData;
-
     // address of staking contract
     address internal stakingContract;
 

contracts/staking/contracts/test/TestProtocolFees.sol

@@ -0,0 +1,41 @@
+/*
+
+  Copyright 2019 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.9;
+pragma experimental ABIEncoderV2;
+
+import "../src/Staking.sol";
+
+
+contract TestProtocolFees is
+    Staking
+{
+    function setPoolIdOfMaker(bytes32 poolId, address makerAddress)
+        external
+    {
+        poolIdByMakerAddress[makerAddress] = poolId;
+    }
+
+    function getActivePoolsByEpoch()
+        external
+        view
+        returns (bytes32[] memory)
+    {
+        return activePoolsThisEpoch;
+    }
+}

contracts/staking/contracts/test/TestProtocolFeesERC20Proxy.sol

@@ -0,0 +1,44 @@
+/*
+
+  Copyright 2019 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.9;
+
+import "@0x/contracts-asset-proxy/contracts/src/ERC20Proxy.sol";
+
+
+contract TestProtocolFeesERC20Proxy is
+    ERC20Proxy
+{
+    event TransferFromCalled(
+        bytes assetData,
+        address from,
+        address to,
+        uint256 amount
+    );
+
+    function transferFrom(
+        bytes calldata assetData,
+        address from,
+        address to,
+        uint256 amount
+    )
+        external
+    {
+        emit TransferFromCalled(assetData, from, to, amount);
+    }
+}

contracts/staking/contracts/test/TestStaking.sol

@@ -0,0 +1,53 @@
+/*
+
+  Copyright 2019 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.9;
+pragma experimental ABIEncoderV2;
+
+import "../src/Staking.sol";
+
+
+contract TestStaking is
+    Staking
+{
+    // Stub out `payProtocolFee` to be the naive payProtocolFee function so that tests will
+    // not fail for WETH protocol fees.
+    function payProtocolFee(
+        address makerAddress,
+        address,
+        uint256
+    )
+        external
+        payable
+        onlyExchange
+    {
+        uint256 amount = msg.value;
+        bytes32 poolId = getStakingPoolIdOfMaker(makerAddress);
+        uint256 _feesCollectedThisEpoch = protocolFeesThisEpochByPool[poolId];
+        protocolFeesThisEpochByPool[poolId] = _feesCollectedThisEpoch.safeAdd(amount);
+        if (_feesCollectedThisEpoch == 0) {
+            activePoolsThisEpoch.push(poolId);
+        }
+    }
+
+    // Stub out `_unwrapWETH` to prevent the calls to `finalizeFees` from failing in tests
+    // that do not relate to protocol fee payments in WETH.
+    function _unwrapWETH()
+        internal
+    {} // solhint-disable-line no-empty-blocks
+}

contracts/staking/src/artifacts.ts

@@ -40,6 +40,9 @@ import * as StakingPoolRewardVault from '../generated-artifacts/StakingPoolRewar
 import * as StakingProxy from '../generated-artifacts/StakingProxy.json';
 import * as TestCobbDouglas from '../generated-artifacts/TestCobbDouglas.json';
 import * as TestLibFixedMath from '../generated-artifacts/TestLibFixedMath.json';
+import * as TestProtocolFees from '../generated-artifacts/TestProtocolFees.json';
+import * as TestProtocolFeesERC20Proxy from '../generated-artifacts/TestProtocolFeesERC20Proxy.json';
+import * as TestStaking from '../generated-artifacts/TestStaking.json';
 import * as TestStorageLayout from '../generated-artifacts/TestStorageLayout.json';
 import * as ZrxVault from '../generated-artifacts/ZrxVault.json';
 export const artifacts = {
@@ -79,5 +82,8 @@ export const artifacts = {
     ZrxVault: ZrxVault as ContractArtifact,
     TestCobbDouglas: TestCobbDouglas as ContractArtifact,
     TestLibFixedMath: TestLibFixedMath as ContractArtifact,
+    TestProtocolFees: TestProtocolFees as ContractArtifact,
+    TestProtocolFeesERC20Proxy: TestProtocolFeesERC20Proxy as ContractArtifact,
+    TestStaking: TestStaking as ContractArtifact,
     TestStorageLayout: TestStorageLayout as ContractArtifact,
 };

contracts/staking/src/wrappers.ts

@@ -38,5 +38,8 @@ export * from '../generated-wrappers/staking_pool_reward_vault';
 export * from '../generated-wrappers/staking_proxy';
 export * from '../generated-wrappers/test_cobb_douglas';
 export * from '../generated-wrappers/test_lib_fixed_math';
+export * from '../generated-wrappers/test_protocol_fees';
+export * from '../generated-wrappers/test_protocol_fees_erc20_proxy';
+export * from '../generated-wrappers/test_staking';
 export * from '../generated-wrappers/test_storage_layout';
 export * from '../generated-wrappers/zrx_vault';

contracts/staking/test/epoch_test.ts

@@ -6,6 +6,8 @@ import { BigNumber } from '@0x/utils';
 import * as chai from 'chai';
 import * as _ from 'lodash';
 
+import { artifacts } from '../src';
+
 import { constants as stakingConstants } from './utils/constants';
 import { StakingWrapper } from './utils/staking_wrapper';
 

contracts/staking/test/protocol_fees.ts

@@ -0,0 +1,361 @@
+import { blockchainTests, constants, expect, LogDecoder } from '@0x/contracts-test-utils';
+import { StakingRevertErrors } from '@0x/order-utils';
+import { BigNumber } from '@0x/utils';
+import { LogWithDecodedArgs } from 'ethereum-types';
+
+import {
+    artifacts,
+    TestProtocolFeesContract,
+    TestProtocolFeesERC20ProxyContract,
+    TestProtocolFeesERC20ProxyTransferFromCalledEventArgs,
+} from '../src';
+
+// tslint:disable:no-unnecessary-type-assertion
+blockchainTests('Protocol Fee Unit Tests', env => {
+    // The accounts that will be used during testing.
+    let owner: string;
+    let exchange: string;
+    let nonExchange: string;
+    let makerAddress: string;
+    let payerAddress: string;
+
+    // The contract that will be used for testng `payProtocolFee` and `_unwrapETH`.
+    let protocolFees: TestProtocolFeesContract;
+    let proxy: TestProtocolFeesERC20ProxyContract;
+
+    // The log decoder that will be used to decode logs from TestProtocolFeesERC20Proxy.
+    let logDecoder: LogDecoder;
+
+    // The default protocol fee that will be paid -- a somewhat realistic value.
+    const DEFAULT_PROTOCOL_FEE_PAID = new BigNumber(150000).times(10000000);
+
+    // The default pool Id that will be used.
+    const DEFAULT_POOL_ID = '0x0000000000000000000000000000000000000000000000000000000000000001';
+
+    // The WETH asset data that should be set in the contract.
+    const WETH_ASSET_DATA = '0xf47261b0000000000000000000000000c02aaa39b223fe8d0a0e5c4f27ead9083c756cc2';
+
+    before(async () => {
+        // Get accounts to represent the exchange and an address that is not a registered exchange.
+        [
+            owner,
+            exchange,
+            nonExchange,
+            makerAddress,
+            payerAddress,
+        ] = (await env.web3Wrapper.getAvailableAddressesAsync()).slice(0, 6);
+
+        // Deploy the protocol fees contract.
+        protocolFees = await TestProtocolFeesContract.deployFrom0xArtifactAsync(
+            artifacts.TestProtocolFees,
+            env.provider,
+            {
+                ...env.txDefaults,
+                from: owner,
+            },
+            {},
+        );
+
+        // Deploy the erc20Proxy for testing.
+        proxy = await TestProtocolFeesERC20ProxyContract.deployFrom0xArtifactAsync(
+            artifacts.TestProtocolFeesERC20Proxy,
+            env.provider,
+            env.txDefaults,
+            {},
+        );
+
+        // Register the test ERC20Proxy in the exchange.
+        await protocolFees.addERC20AssetProxy.awaitTransactionSuccessAsync(proxy.address);
+
+        // Register an exchange in the protocol fee contract.
+        await protocolFees.addExchangeAddress.awaitTransactionSuccessAsync(exchange, { from: owner });
+
+        // "Register" the makerAddress in the default pool.
+        await protocolFees.setPoolIdOfMaker.awaitTransactionSuccessAsync(DEFAULT_POOL_ID, makerAddress);
+
+        // Create the log decoder that will be used to decode TransferFromCalledEvent logs.
+        logDecoder = new LogDecoder(env.web3Wrapper, artifacts);
+    });
+
+    blockchainTests.resets('payProtocolFee', () => {
+        // Verify that the DEFAULT_POOL_ID was pushed to the active pool list and that the correct amount
+        // is registered in the pool, or that the NIL_POOL's state was unaffected depending on which pool id
+        // was provided.
+        async function verifyEndStateAsync(poolId: string, amount: BigNumber): Promise<void> {
+            if (poolId === DEFAULT_POOL_ID) {
+                // Ensure that the `DEFAULT_POOL_ID` was pushed into this epoch's active pool.
+                const activePools = await protocolFees.getActivePoolsByEpoch.callAsync();
+                expect(activePools.length).to.be.eq(1);
+                expect(activePools[0]).to.be.eq(DEFAULT_POOL_ID);
+
+                // Ensure that the `DEFAULT_PROTOCOL_FEE_PAID` was attributed to the maker's pool.
+                const feesInMakerPool = await protocolFees.getProtocolFeesThisEpochByPool.callAsync(DEFAULT_POOL_ID);
+                expect(feesInMakerPool).bignumber.to.be.eq(amount);
+            } else {
+                // Ensure that the only active pool this epoch is the "zero" pool.
+                const activePools = await protocolFees.getActivePoolsByEpoch.callAsync();
+                expect(activePools.length).to.be.eq(0);
+
+                // Ensure that the `NIL_POOL` was not attributed a payment.
+                const feesInMakerPool = await protocolFees.getProtocolFeesThisEpochByPool.callAsync(
+                    constants.NULL_BYTES32,
+                );
+                expect(feesInMakerPool).bignumber.to.be.eq(constants.ZERO_AMOUNT);
+            }
+        }
+
+        it('should revert if called by a non-exchange', async () => {
+            const expectedError = new StakingRevertErrors.OnlyCallableByExchangeError(nonExchange);
+            const tx = protocolFees.payProtocolFee.sendTransactionAsync(
+                makerAddress,
+                payerAddress,
+                DEFAULT_PROTOCOL_FEE_PAID,
+                { from: nonExchange },
+            );
+            return expect(tx).to.revertWith(expectedError);
+        });
+
+        it('should revert if `protocolFeePaid` is zero with zero value sent', async () => {
+            const expectedError = new StakingRevertErrors.InvalidProtocolFeePaymentError(
+                constants.ZERO_AMOUNT,
+                constants.ZERO_AMOUNT,
+            );
+            const tx = protocolFees.payProtocolFee.sendTransactionAsync(
+                makerAddress,
+                payerAddress,
+                constants.ZERO_AMOUNT,
+                { from: exchange, value: constants.ZERO_AMOUNT },
+            );
+            return expect(tx).to.revertWith(expectedError);
+        });
+
+        it('should revert if `protocolFeePaid` is zero with non-zero value sent', async () => {
+            const expectedError = new StakingRevertErrors.InvalidProtocolFeePaymentError(
+                constants.ZERO_AMOUNT,
+                DEFAULT_PROTOCOL_FEE_PAID,
+            );
+            const tx = protocolFees.payProtocolFee.sendTransactionAsync(
+                makerAddress,
+                payerAddress,
+                constants.ZERO_AMOUNT,
+                { from: exchange, value: DEFAULT_PROTOCOL_FEE_PAID },
+            );
+            return expect(tx).to.revertWith(expectedError);
+        });
+
+        it('should revert if `protocolFeePaid` is different than the provided message value', async () => {
+            const differentProtocolFeePaid = DEFAULT_PROTOCOL_FEE_PAID.minus(50);
+            const expectedError = new StakingRevertErrors.InvalidProtocolFeePaymentError(
+                differentProtocolFeePaid,
+                DEFAULT_PROTOCOL_FEE_PAID,
+            );
+            const tx = protocolFees.payProtocolFee.sendTransactionAsync(
+                makerAddress,
+                payerAddress,
+                differentProtocolFeePaid,
+                { from: exchange, value: DEFAULT_PROTOCOL_FEE_PAID },
+            );
+            return expect(tx).to.revertWith(expectedError);
+        });
+
+        it('should call `transferFrom` in the proxy if no value is sent and the maker is not in a pool', async () => {
+            const receipt = await protocolFees.payProtocolFee.awaitTransactionSuccessAsync(
+                payerAddress, // This is an unregistered maker address
+                payerAddress,
+                DEFAULT_PROTOCOL_FEE_PAID,
+                { from: exchange, value: 0 },
+            );
+
+            // Ensure that the correct number of logs were recorded.
+            expect(receipt.logs.length).to.be.eq(1);
+
+            // Ensure that the correct log was recorded.
+            const log = logDecoder.decodeLogOrThrow(receipt.logs[0]) as LogWithDecodedArgs<
+                TestProtocolFeesERC20ProxyTransferFromCalledEventArgs
+            >;
+            expect(log.event).to.be.eq('TransferFromCalled');
+            expect(log.args.assetData).to.be.eq(WETH_ASSET_DATA);
+            expect(log.args.amount).bignumber.to.be.eq(DEFAULT_PROTOCOL_FEE_PAID);
+            expect(log.args.from).to.be.eq(payerAddress);
+            expect(log.args.to).to.be.eq(protocolFees.address);
+
+            // Verify that the end state is correct.
+            await verifyEndStateAsync(constants.NULL_BYTES32, constants.ZERO_AMOUNT);
+        });
+
+        it('should call `transferFrom` in the proxy and update `protocolFeesThisEpochByPool` if no value is sent and the maker is in a pool', async () => {
+            const receipt = await protocolFees.payProtocolFee.awaitTransactionSuccessAsync(
+                makerAddress,
+                payerAddress,
+                DEFAULT_PROTOCOL_FEE_PAID,
+                { from: exchange, value: 0 },
+            );
+
+            // Ensure that the correct number of logs were recorded.
+            expect(receipt.logs.length).to.be.eq(1);
+
+            // Ensure that the correct log was recorded.
+            const log = logDecoder.decodeLogOrThrow(receipt.logs[0]) as LogWithDecodedArgs<
+                TestProtocolFeesERC20ProxyTransferFromCalledEventArgs
+            >;
+            expect(log.event).to.be.eq('TransferFromCalled');
+            expect(log.args.assetData).to.be.eq(WETH_ASSET_DATA);
+            expect(log.args.amount).bignumber.to.be.eq(DEFAULT_PROTOCOL_FEE_PAID);
+            expect(log.args.from).to.be.eq(payerAddress);
+            expect(log.args.to).to.be.eq(protocolFees.address);
+
+            // Verify that the end state is correct.
+            await verifyEndStateAsync(DEFAULT_POOL_ID, DEFAULT_PROTOCOL_FEE_PAID);
+        });
+
+        it('should not call `transferFrom` in the proxy and should not update `protocolFeesThisEpochByPool` if value is sent and the maker is not in a pool', async () => {
+            const receipt = await protocolFees.payProtocolFee.awaitTransactionSuccessAsync(
+                payerAddress, // This is an unregistered maker address
+                payerAddress,
+                DEFAULT_PROTOCOL_FEE_PAID,
+                { from: exchange, value: DEFAULT_PROTOCOL_FEE_PAID },
+            );
+
+            // Ensure that the correct number of logs were recorded.
+            expect(receipt.logs.length).to.be.eq(0);
+
+            // Verify that the end state is correct.
+            await verifyEndStateAsync(constants.NULL_BYTES32, constants.ZERO_AMOUNT);
+        });
+
+        it('should not call `transferFrom` in the proxy and should update `protocolFeesThisEpochByPool` if value is sent and the maker is in a pool', async () => {
+            const receipt = await protocolFees.payProtocolFee.awaitTransactionSuccessAsync(
+                makerAddress,
+                payerAddress,
+                DEFAULT_PROTOCOL_FEE_PAID,
+                { from: exchange, value: DEFAULT_PROTOCOL_FEE_PAID },
+            );
+
+            // Ensure that the correct number of logs were recorded.
+            expect(receipt.logs.length).to.be.eq(0);
+
+            // Verify that the end state is correct.
+            await verifyEndStateAsync(DEFAULT_POOL_ID, DEFAULT_PROTOCOL_FEE_PAID);
+        });
+
+        it('should only have one active pool if a fee is paid on behalf of one maker ETH twice', async () => {
+            // Pay the first fee
+            await protocolFees.payProtocolFee.awaitTransactionSuccessAsync(
+                makerAddress,
+                payerAddress,
+                DEFAULT_PROTOCOL_FEE_PAID,
+                { from: exchange, value: DEFAULT_PROTOCOL_FEE_PAID },
+            );
+
+            // Pay the second fee
+            const receipt = await protocolFees.payProtocolFee.awaitTransactionSuccessAsync(
+                makerAddress,
+                payerAddress,
+                DEFAULT_PROTOCOL_FEE_PAID,
+                { from: exchange, value: DEFAULT_PROTOCOL_FEE_PAID },
+            );
+
+            // Ensure that the correct number of logs were recorded.
+            expect(receipt.logs.length).to.be.eq(0);
+
+            // Verify that the end state is correct -- namely that the active pools list was only updated once,
+            // and that the correct amount is recorded on behalf of the maker.
+            await verifyEndStateAsync(DEFAULT_POOL_ID, DEFAULT_PROTOCOL_FEE_PAID.times(2));
+        });
+
+        it('should only have one active pool if a fee is paid on behalf of one maker in WETH and then ETH', async () => {
+            // Pay the first fee
+            await protocolFees.payProtocolFee.awaitTransactionSuccessAsync(
+                makerAddress,
+                payerAddress,
+                DEFAULT_PROTOCOL_FEE_PAID,
+                { from: exchange, value: 0 },
+            );
+
+            // Pay the second fee
+            const receipt = await protocolFees.payProtocolFee.awaitTransactionSuccessAsync(
+                makerAddress,
+                payerAddress,
+                DEFAULT_PROTOCOL_FEE_PAID,
+                { from: exchange, value: DEFAULT_PROTOCOL_FEE_PAID },
+            );
+
+            // Ensure that the correct number of logs were recorded.
+            expect(receipt.logs.length).to.be.eq(0);
+
+            // Verify that the end state is correct -- namely that the active pools list was only updated once,
+            // and that the correct amount is recorded on behalf of the maker.
+            await verifyEndStateAsync(DEFAULT_POOL_ID, DEFAULT_PROTOCOL_FEE_PAID.times(2));
+        });
+
+        it('should only have one active pool if a fee is paid on behalf of one maker in ETH and then WETH', async () => {
+            // Pay the first fee
+            await protocolFees.payProtocolFee.awaitTransactionSuccessAsync(
+                makerAddress,
+                payerAddress,
+                DEFAULT_PROTOCOL_FEE_PAID,
+                { from: exchange, value: DEFAULT_PROTOCOL_FEE_PAID },
+            );
+
+            // Pay the second fee
+            const receipt = await protocolFees.payProtocolFee.awaitTransactionSuccessAsync(
+                makerAddress,
+                payerAddress,
+                DEFAULT_PROTOCOL_FEE_PAID,
+                { from: exchange, value: 0 },
+            );
+
+            // Ensure that the correct number of logs were recorded.
+            expect(receipt.logs.length).to.be.eq(1);
+
+            // Ensure that the correct log was recorded
+            const log = logDecoder.decodeLogOrThrow(receipt.logs[0]) as LogWithDecodedArgs<
+                TestProtocolFeesERC20ProxyTransferFromCalledEventArgs
+            >;
+            expect(log.event).to.be.eq('TransferFromCalled');
+            expect(log.args.assetData).to.be.eq(WETH_ASSET_DATA);
+            expect(log.args.amount).bignumber.to.be.eq(DEFAULT_PROTOCOL_FEE_PAID);
+            expect(log.args.from).to.be.eq(payerAddress);
+            expect(log.args.to).to.be.eq(protocolFees.address);
+
+            // Verify that the end state is correct -- namely that the active pools list was only updated once,
+            // and that the correct amount is recorded on behalf of the maker.
+            await verifyEndStateAsync(DEFAULT_POOL_ID, DEFAULT_PROTOCOL_FEE_PAID.times(2));
+        });
+
+        it('should only have one active pool if a fee is paid on behalf of one maker in WETH twice', async () => {
+            // Pay the first fee
+            await protocolFees.payProtocolFee.awaitTransactionSuccessAsync(
+                makerAddress,
+                payerAddress,
+                DEFAULT_PROTOCOL_FEE_PAID,
+                { from: exchange, value: 0 },
+            );
+
+            // Pay the second fee
+            const receipt = await protocolFees.payProtocolFee.awaitTransactionSuccessAsync(
+                makerAddress,
+                payerAddress,
+                DEFAULT_PROTOCOL_FEE_PAID,
+                { from: exchange, value: 0 },
+            );
+
+            // Ensure that the correct number of logs were recorded.
+            expect(receipt.logs.length).to.be.eq(1);
+
+            // Ensure that the correct log was recorded.
+            const log = logDecoder.decodeLogOrThrow(receipt.logs[0]) as LogWithDecodedArgs<
+                TestProtocolFeesERC20ProxyTransferFromCalledEventArgs
+            >;
+            expect(log.event).to.be.eq('TransferFromCalled');
+            expect(log.args.assetData).to.be.eq(WETH_ASSET_DATA);
+            expect(log.args.amount).bignumber.to.be.eq(DEFAULT_PROTOCOL_FEE_PAID);
+            expect(log.args.from).to.be.eq(payerAddress);
+            expect(log.args.to).to.be.eq(protocolFees.address);
+
+            // Verify that the end state is correct -- namely that the active pools list was only updated once,
+            // and that the correct amount is recorded on behalf of the maker.
+            await verifyEndStateAsync(DEFAULT_POOL_ID, DEFAULT_PROTOCOL_FEE_PAID.times(2));
+        });
+    });
+});

contracts/staking/test/rewards_test.ts

@@ -4,6 +4,8 @@ import { blockchainTests, describe, expect, provider, web3Wrapper } from '@0x/co
 import { BigNumber } from '@0x/utils';
 import * as _ from 'lodash';
 
+import { artifacts } from '../src';
+
 import { FinalizerActor } from './actors/finalizer_actor';
 import { StakerActor } from './actors/staker_actor';
 import { StakingWrapper } from './utils/staking_wrapper';

contracts/staking/test/stake_test.ts

@@ -5,6 +5,8 @@ import { StakingRevertErrors } from '@0x/order-utils';
 import { BigNumber } from '@0x/utils';
 import * as _ from 'lodash';
 
+import { artifacts } from '../src';
+
 import { StakerActor } from './actors/staker_actor';
 import { StakingWrapper } from './utils/staking_wrapper';
 import { StakeInfo, StakeStatus } from './utils/types';

contracts/staking/test/utils/staking_wrapper.ts

@@ -62,7 +62,7 @@ export class StakingWrapper {
     constructor(
         provider: Provider,
         ownerAddres: string,
-        erc20ProxyContract: ERC20ProxyContract,
+        erc20ProxyContract: any, // This needs to be the `any` type so that other types of proxies can be used
         zrxTokenContract: DummyERC20TokenContract,
     ) {
         this._web3Wrapper = new Web3Wrapper(provider);

contracts/staking/tsconfig.json

@@ -38,6 +38,9 @@
         "generated-artifacts/StakingProxy.json",
         "generated-artifacts/TestCobbDouglas.json",
         "generated-artifacts/TestLibFixedMath.json",
+        "generated-artifacts/TestProtocolFees.json",
+        "generated-artifacts/TestProtocolFeesERC20Proxy.json",
+        "generated-artifacts/TestStaking.json",
         "generated-artifacts/TestStorageLayout.json",
         "generated-artifacts/ZrxVault.json"
     ],

packages/order-utils/src/staking_revert_errors.ts

@@ -179,6 +179,19 @@ export class InvalidStakeStatusError extends RevertError {
     }
 }
 
+export class InvalidProtocolFeePaymentError extends RevertError {
+    constructor(
+        expectedProtocolFeePaid?: BigNumber | number | string,
+        actualProtocolFeePaid?: BigNumber | number | string,
+    ) {
+        super(
+            'InvalidProtocolFeePaymentError',
+            'InvalidProtocolFeePaymentError(uint256 expectedProtocolFeePaid, uint256 actualProtocolFeePaid)',
+            { expectedProtocolFeePaid, actualProtocolFeePaid },
+        );
+    }
+}
+
 const types = [
     MiscalculatedRewardsError,
     OnlyCallableByExchangeError,
@@ -200,6 +213,7 @@ const types = [
     EthVaultNotSetError,
     RewardVaultNotSetError,
     InvalidStakeStatusError,
+    InvalidProtocolFeePaymentError,
 ];
 
 // Register the types we've defined.