Batch fetch balances protocol viewer (#540)

contracts/viewer/lib/ERC20Viewer.sol

@@ -56,6 +56,34 @@ contract ERC20Viewer {
         return balances;
     }
 
+    /*
+     * Fetches token balances for each tokenAddress, tokenOwner pair
+     *
+     * @param  _tokenAddresses    Addresses of ERC20 contracts to check balance for
+     * @param  _tokenOwners       Addresses of users sequential to tokenAddress to fetch balance for
+     * @return  uint256[]         Array of balances for each ERC20 contract passed in
+     */
+    function batchFetchUsersBalances(
+        address[] calldata _tokenAddresses,
+        address[] calldata _tokenOwners
+    )
+        external
+        returns (uint256[] memory)
+    {
+        // Cache length of addresses to fetch balances for
+        uint256 _addressesCount = _tokenAddresses.length;
+
+        // Instantiate output array in memory
+        uint256[] memory balances = new uint256[](_addressesCount);
+
+        // Cycle through contract addresses array and fetching the balance of each for the owner
+        for (uint256 i = 0; i < _addressesCount; i++) {
+            balances[i] = IERC20(_tokenAddresses[i]).balanceOf(_tokenOwners[i]);
+        }
+
+        return balances;
+    }
+
     /*
      * Fetches multiple supplies for passed in array of ERC20 contract addresses
      *

contracts/viewer/lib/RebalancingSetTokenViewer.sol

@@ -96,4 +96,30 @@ contract RebalancingSetTokenViewer {
 
         return (rebalanceState, auctionIntegerParams);
     }
+
+    /*
+     * Fetches RebalancingSetToken states for an array of RebalancingSetToken instances
+     *
+     * @param  _rebalancingSetTokens[]       RebalancingSetToken contract instances
+     * @return RebalancingLibrary.State[]    Current rebalance states on the RebalancingSetToken
+     */
+    function batchFetchRebalanceStateAsync(
+        IRebalancingSetToken[] calldata _rebalancingSetTokens
+    )
+        external
+        returns (RebalancingLibrary.State[] memory)
+    {
+        // Cache length of addresses to fetch states for
+        uint256 _addressesCount = _rebalancingSetTokens.length;
+
+        // Instantiate output array in memory
+        RebalancingLibrary.State[] memory states = new RebalancingLibrary.State[](_addressesCount);
+
+        // Cycle through contract addresses array and fetching the rebalance state of each RebalancingSet
+        for (uint256 i = 0; i < _addressesCount; i++) {
+            states[i] = _rebalancingSetTokens[i].rebalanceState();
+        }
+
+        return states;
+    }
 }

test/contracts/viewer/lib/erc20Viewer.spec.ts

@@ -2,8 +2,8 @@ require('module-alias/register');
 
 import * as ABIDecoder from 'abi-decoder';
 import * as chai from 'chai';
-import { BigNumber } from 'bignumber.js';
 import { Address } from 'set-protocol-utils';
+import { BigNumber } from 'bignumber.js';
 
 import ChaiSetup from '@utils/chaiSetup';
 import { BigNumberSetup } from '@utils/bigNumberSetup';
@@ -47,6 +47,7 @@ contract('ERC20Viewer', accounts => {
     deployerAccount,
     managerAccount,
     ownerAccount,
+    anotherAccount,
   ] = accounts;
 
   let coreMock: CoreMockContract;
@@ -198,13 +199,13 @@ contract('ERC20Viewer', accounts => {
       const balances: BigNumber[] = await subject();
       const balancesJSON = JSON.stringify(balances);
 
-      const expectedSupplies = await erc20Helper.getTokenBalances(
+      const expectedBalances = await erc20Helper.getTokenBalances(
         [token, currentSetToken, rebalancingSetToken],
         subjectTokenOwner
       );
-      const expectedSuppliesJSON = JSON.stringify(expectedSupplies);
+      const expectedBalancesJSON = JSON.stringify(expectedBalances);
 
-      expect(balancesJSON).to.equal(expectedSuppliesJSON);
+      expect(balancesJSON).to.equal(expectedBalancesJSON);
     });
 
     describe('when the token addresses includes a non ERC20 contract', async () => {
@@ -217,4 +218,44 @@ contract('ERC20Viewer', accounts => {
       });
     });
   });
+
+  describe('#batchFetchUsersBalances', async () => {
+    let subjectTokenAddresses: Address[];
+    let subjectTokenOwners: Address[];
+
+    beforeEach(async () => {
+      token = await erc20Helper.deployTokenAsync(ownerAccount);
+
+      subjectTokenAddresses = [token.address, currentSetToken.address, rebalancingSetToken.address];
+      subjectTokenOwners = [deployerAccount, anotherAccount, deployerAccount];
+    });
+
+    async function subject(): Promise<BigNumber[]> {
+      return erc20Viewer.batchFetchUsersBalances.callAsync(
+        subjectTokenAddresses,
+        subjectTokenOwners,
+      );
+    }
+
+    it('fetches the balances of the token addresses', async () => {
+      const balances: BigNumber[] = await subject();
+      const balancesJSON = JSON.stringify(balances);
+
+      const expectedDeployerBalances = await erc20Helper.getTokenBalances(
+        [token, rebalancingSetToken],
+        deployerAccount
+      );
+      const expectedAnotherAccountBalances = await erc20Helper.getTokenBalances(
+        [currentSetToken],
+        anotherAccount
+      );
+      const expectedBalances = [
+        expectedDeployerBalances[0],
+        expectedAnotherAccountBalances[0],
+        expectedDeployerBalances[1],
+      ];
+      const expectedBalancesJSON = JSON.stringify(expectedBalances);
+      expect(balancesJSON).to.equal(expectedBalancesJSON);
+    });
+  });
 });

test/contracts/viewer/lib/rebalancingSetTokenViewer.spec.ts

@@ -4,6 +4,7 @@ import * as ABIDecoder from 'abi-decoder';
 import * as chai from 'chai';
 import * as setProtocolUtils from 'set-protocol-utils';
 import { Address } from 'set-protocol-utils';
+import { BigNumber } from 'bignumber.js';
 
 import ChaiSetup from '@utils/chaiSetup';
 import { BigNumberSetup } from '@utils/bigNumberSetup';
@@ -372,4 +373,81 @@ contract('ProtocolViewer', accounts => {
       });
     });
   });
+
+  describe.only('#batchFetchRebalanceStateAsync', async () => {
+    let subjectRebalancingSetAddresses: Address[];
+
+    let rebalancingSetToken: RebalancingSetTokenContract;
+    let currentSetToken: SetTokenContract;
+    let nextSetToken: SetTokenContract;
+
+    let defaultRebalancingSetToken: RebalancingSetTokenContract;
+
+    beforeEach(async () => {
+      const naturalUnits = [ether(.001), ether(.0001)];
+
+      const setTokens = await rebalancingHelper.createSetTokensAsync(
+        coreMock,
+        factory.address,
+        transferProxy.address,
+        2,
+        naturalUnits
+      );
+
+      currentSetToken = setTokens[0];
+      nextSetToken = setTokens[1];
+
+      rebalancingSetToken = await rebalancingHelper.createDefaultRebalancingSetTokenAsync(
+        coreMock,
+        rebalancingFactory.address,
+        managerAccount,
+        currentSetToken.address,
+        ONE_DAY_IN_SECONDS
+      );
+
+      // Issue currentSetToken
+      await coreMock.issue.sendTransactionAsync(currentSetToken.address, ether(8), {from: deployerAccount});
+      await erc20Helper.approveTransfersAsync([currentSetToken], transferProxy.address);
+
+      // Use issued currentSetToken to issue rebalancingSetToken
+      const rebalancingSetTokenQuantityToIssue = ether(8);
+      await coreMock.issue.sendTransactionAsync(rebalancingSetToken.address, rebalancingSetTokenQuantityToIssue);
+
+      // Transition original rebalancing set to proposal
+      await rebalancingHelper.defaultTransitionToProposeAsync(
+          coreMock,
+          rebalancingComponentWhiteList,
+          rebalancingSetToken,
+          nextSetToken,
+          constantAuctionPriceCurve.address,
+          managerAccount
+        );
+
+      defaultRebalancingSetToken = await rebalancingHelper.createDefaultRebalancingSetTokenAsync(
+        coreMock,
+        rebalancingFactory.address,
+        managerAccount,
+        currentSetToken.address,
+        ONE_DAY_IN_SECONDS
+      );
+
+      subjectRebalancingSetAddresses = [rebalancingSetToken.address, defaultRebalancingSetToken.address];
+    });
+
+    async function subject(): Promise<BigNumber[]> {
+      return rebalancingSetTokenViewer.batchFetchRebalanceStateAsync.callAsync(
+        subjectRebalancingSetAddresses,
+      );
+    }
+
+    it('fetches the RebalancingSetTokens\' states', async () => {
+      const rebalanceAuctionStates: BigNumber[] = await subject();
+
+      const firstRebalancingSetState = rebalanceAuctionStates[0];
+      expect(firstRebalancingSetState).to.be.bignumber.equal(SetUtils.REBALANCING_STATE.PROPOSAL);
+
+      const secondRebalancingSetState = rebalanceAuctionStates[1];
+      expect(secondRebalancingSetState).to.be.bignumber.equal(SetUtils.REBALANCING_STATE.DEFAULT);
+    });
+  });
 });