From 457b2a108ae84a9bbe5f757ea18f2b3d4ff2dacb Mon Sep 17 00:00:00 2001
From: alexcos20 <alex.coseru@gmail.com>
Date: Wed, 14 Oct 2020 02:57:45 -0700
Subject: [PATCH] add getTokensRemovedforPoolShares

---
 src/balancer/OceanPool.ts           | 31 +++++++++++++++++++++++++++++
 test/unit/balancer/Balancer.test.ts | 15 ++++----------
 2 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/src/balancer/OceanPool.ts b/src/balancer/OceanPool.ts
index 3153c0a1..2a08b5df 100644
--- a/src/balancer/OceanPool.ts
+++ b/src/balancer/OceanPool.ts
@@ -5,6 +5,7 @@ import { Pool } from './Pool'
 import { EventData, Filter } from 'web3-eth-contract'
 import BigNumber from 'bignumber.js'
 import Decimal from 'decimal.js'
+import { createDiffieHellman } from 'crypto'
 
 declare type PoolTransactionType = 'swap' | 'join' | 'exit'
 
@@ -13,6 +14,11 @@ export interface PoolDetails {
   tokens: string[]
 }
 
+export interface TokensReceived {
+  dtAmount: string
+  oceanAmount: string
+}
+
 export interface PoolTransaction {
   poolAddress: string
   dtAddress: string
@@ -385,6 +391,31 @@ export class OceanPool extends Pool {
     return this.calcSingleOutGivenPoolIn(poolAddress, this.oceanAddress, poolShares)
   }
 
+  /**
+   * Returns datatoken & Ocean  amounts received after spending poolShares
+   * @param poolAddress
+   * @param poolShares
+   * @return {TokensReceived}
+   */
+  public async getTokensRemovedforPoolShares(
+    poolAddress: string,
+    poolShares: string
+  ): Promise<TokensReceived> {
+    const amounts = { dtAmount: '0', oceanAmount: '0' }
+    try {
+      const totalPoolTokens = await this.getPoolSharesTotalSupply(poolAddress)
+      amounts.dtAmount = String(
+        (Number(poolShares) / Number(totalPoolTokens)) *
+          Number(await this.getDTReserve(poolAddress))
+      )
+      amounts.oceanAmount = String(
+        (Number(poolShares) / Number(totalPoolTokens)) *
+          Number(await this.getOceanReserve(poolAddress))
+      )
+    } catch (e) {}
+    return amounts
+  }
+
   /**
    * Returns max DT amount that you can add to the pool
    * @param poolAddress
diff --git a/test/unit/balancer/Balancer.test.ts b/test/unit/balancer/Balancer.test.ts
index d574ac80..ec866386 100644
--- a/test/unit/balancer/Balancer.test.ts
+++ b/test/unit/balancer/Balancer.test.ts
@@ -388,18 +388,11 @@ describe('Balancer flow', () => {
     assert(parseFloat(bobDtBalance) < parseFloat(newbobDtBalance))
     assert(parseFloat(poolShares) > parseFloat(newpoolShares))
   })
-
-  it('ALice should remove all liquidity', async () => {
+  it('ALice should know how many tokens she will get for removing all liquidity', async () => {
     const aliceShares = await Pool.sharesBalance(alice, greatPool)
-    const aliceDtBalance = await datatoken.balance(tokenAddress, alice)
-    const aliceOceanBalance = await datatoken.balance(oceanTokenAddress, alice)
-    await Pool.removePoolLiquidity(alice, greatPool, aliceShares)
-    const newAliceDtBalance = await datatoken.balance(tokenAddress, alice)
-    const newAliceOceanBalance = await datatoken.balance(oceanTokenAddress, alice)
-    const newAliceShares = await Pool.sharesBalance(alice, greatPool)
-    assert(parseFloat(aliceDtBalance) < parseFloat(newAliceDtBalance))
-    assert(parseFloat(aliceOceanBalance) < parseFloat(newAliceOceanBalance))
-    assert(parseFloat(aliceShares) > parseFloat(newAliceShares))
+    const amounts = await Pool.getTokensRemovedforPoolShares(greatPool, aliceShares)
+    assert(parseFloat(amounts.dtAmount) > 0)
+    assert(parseFloat(amounts.oceanAmount) > 0)
   })
   it('ALice should get all the pools that she created', async () => {
     const alicePools = await Pool.getPoolsbyCreator(alice)