From ee926aca132b284051db53c40ab5e53ed439c255 Mon Sep 17 00:00:00 2001
From: Felix Leupold <felixleupold90@gmail.com>
Date: Wed, 12 Feb 2020 19:12:50 +0100
Subject: [PATCH] [BatchExchangeViewer] allow to fetch balances

---
 contracts/BatchExchangeViewer.sol     | 37 +++++++++++++++++
 test/stablex/batch_exchange_viewer.js | 58 +++++++++++++++++++++++++++
 2 files changed, 95 insertions(+)

diff --git a/contracts/BatchExchangeViewer.sol b/contracts/BatchExchangeViewer.sol
index 249fd14c9..8ba65a39a 100644
--- a/contracts/BatchExchangeViewer.sol
+++ b/contracts/BatchExchangeViewer.sol
@@ -1,6 +1,8 @@
 pragma solidity ^0.5.0;
+pragma experimental ABIEncoderV2;
 
 import "solidity-bytes-utils/contracts/BytesLib.sol";
+import "openzeppelin-solidity/contracts/math/Math.sol";
 import "./BatchExchange.sol";
 
 
@@ -16,6 +18,41 @@ contract BatchExchangeViewer {
         batchExchange = exchange;
     }
 
+    struct Balance {
+        address token; // The token for which this balance is reported
+        uint256 availableSellAmount; // Amount available to trade in the auction for which it was requested
+        uint256 pendingDeposit; // Amount that was deposited in the requested auction. Will be tradeable in the next auction
+        uint256 withdrawableBalance; // Amount that can be withdrawn in the requested auction
+    }
+
+    /** @dev Returns the user's token balances for the auction that is currently being settled
+     */
+    function getBalances(address user) public view returns (Balance[] memory) {
+        uint32 currentBatch = batchExchange.getCurrentBatchId();
+        Balance[] memory result = new Balance[](batchExchange.numTokens());
+        for (uint16 index = 0; index < result.length; index++) {
+            address token = batchExchange.tokenIdToAddressMap(index);
+
+            (uint256 depositBalance, uint32 depositedInBatch) = batchExchange.getPendingDeposit(user, token);
+            if (depositedInBatch < currentBatch) {
+                depositBalance = 0;
+            }
+
+            (uint256 withdrawableBalance, uint32 withdrawRequestedInBatch) = batchExchange.getPendingWithdraw(user, token);
+            if (withdrawRequestedInBatch >= currentBatch) {
+                withdrawableBalance = 0;
+            }
+
+            result[index] = Balance({
+                token: token,
+                availableSellAmount: batchExchange.getBalance(user, token),
+                pendingDeposit: depositBalance,
+                withdrawableBalance: withdrawableBalance
+            });
+        }
+        return result;
+    }
+
     /** @dev Queries the orderbook for the auction that is still accepting orders
      *  @param tokenFilter all returned order will have buy *and* sell token from this list (leave empty for "no filter")
      *  @return encoded bytes representing orders
diff --git a/test/stablex/batch_exchange_viewer.js b/test/stablex/batch_exchange_viewer.js
index 2170c46d8..cb48277c1 100644
--- a/test/stablex/batch_exchange_viewer.js
+++ b/test/stablex/batch_exchange_viewer.js
@@ -20,6 +20,64 @@ contract("BatchExchangeViewer", accounts => {
     await batchExchange.addToken(token_2.address)
   })
 
+  describe("getBalances", () => {
+    it("returns all listed tokens", async () => {
+      const viewer = await BatchExchangeViewer.new(batchExchange.address)
+      const result = await viewer.getBalances(accounts[0])
+      assert.equal(result.filter(balance => balance.availableSellAmount == 0).length, 3)
+    })
+
+    it("returns pending deposits", async () => {
+      await token_1.givenAnyReturnBool(true)
+      await batchExchange.deposit(token_1.address, 100)
+
+      const viewer = await BatchExchangeViewer.new(batchExchange.address)
+      const result = await viewer.getBalances(accounts[0])
+      assert.equal(result[1].pendingDeposit, 100)
+    })
+
+    it("returns available balance once deposit is processed", async () => {
+      await token_1.givenAnyReturnBool(true)
+      await batchExchange.deposit(token_1.address, 100)
+      await closeAuction(batchExchange)
+
+      const viewer = await BatchExchangeViewer.new(batchExchange.address)
+      const result = await viewer.getBalances(accounts[0])
+      assert.equal(result[1].pendingDeposit, 0)
+      assert.equal(result[1].availableSellAmount, 100)
+    })
+
+    it("returns withdrawable balance once it is claimable", async () => {
+      await token_1.givenAnyReturnBool(true)
+      await batchExchange.deposit(token_1.address, 100)
+      await closeAuction(batchExchange)
+      await batchExchange.requestWithdraw(token_1.address, 100)
+
+      const viewer = await BatchExchangeViewer.new(batchExchange.address)
+      const before_valid = await viewer.getBalances(accounts[0])
+      assert.equal(before_valid[1].availableSellAmount, 100)
+      assert.equal(before_valid[1].withdrawableBalance, 0)
+
+      await closeAuction(batchExchange)
+
+      const after_valid = await viewer.getBalances(accounts[0])
+      assert.equal(after_valid[1].availableSellAmount, 0)
+      assert.equal(after_valid[1].withdrawableBalance, 100)
+    })
+
+    it("reports withdrawable balance higher than available balance", async () => {
+      // Due to https://github.com/gnosis/dex-contracts/issues/539
+      await token_1.givenAnyReturnBool(true)
+      await batchExchange.deposit(token_1.address, 100)
+      await batchExchange.requestWithdraw(token_1.address, 200)
+      await closeAuction(batchExchange)
+
+      const viewer = await BatchExchangeViewer.new(batchExchange.address)
+      const result = await viewer.getBalances(accounts[0])
+      assert.equal(result[1].withdrawableBalance, 200)
+    })
+  })
+
   describe("getOpenOrderBook", () => {
     it("can be queried without pagination", async () => {
       const batchId = await batchExchange.getCurrentBatchId()