Fix available balance when buying

4 files changed, 250 insertions, 30 deletions

diff --git a/market.py b/market.py
index 9550b77..d7b05ce 100644
--- a/market.py
+++ b/market.py
@@ -220,23 +220,20 @@ class Market:
                 compute_value=compute_value, only=only,
                 repartition=repartition, available_balance_only=available_balance_only)
 
-        values_in_base = self.balances.in_currency(base_currency,
-                compute_value=compute_value)
         if available_balance_only:
-            balance = self.balances.all.get(base_currency)
-            if balance is None:
-                total_base_value = portfolio.Amount(base_currency, 0)
-            else:
-                total_base_value = balance.exchange_free + balance.margin_available
+            repartition, total_base_value, values_in_base = self.balances.available_balances_for_repartition(
+                    base_currency=base_currency, liquidity=liquidity,
+                    repartition=repartition, compute_value=compute_value)
         else:
+            values_in_base = self.balances.in_currency(base_currency,
+                    compute_value=compute_value)
             total_base_value = sum(values_in_base.values())
         new_repartition = self.balances.dispatch_assets(total_base_value,
                 liquidity=liquidity, repartition=repartition)
         if available_balance_only:
             for currency, amount in values_in_base.items():
-                if currency != base_currency:
-                    new_repartition.setdefault(currency, portfolio.Amount(base_currency, 0))
-                    new_repartition[currency] += amount
+                if currency != base_currency and currency not in new_repartition:
+                    new_repartition[currency] = amount
 
         self.trades.compute_trades(values_in_base, new_repartition, only=only)
 
diff --git a/store.py b/store.py
index 32c4121..1a1ed76 100644
--- a/store.py
+++ b/store.py
@@ -325,6 +325,49 @@ class BalanceStore:
         else:
             self.market.report.log_balances(tag=tag, checkpoint=checkpoint)
 
+    def available_balances_for_repartition(self,
+            compute_value="average", base_currency="BTC",
+            liquidity="medium", repartition=None):
+        if repartition is None:
+            repartition = Portfolio.repartition(liquidity=liquidity)
+        base_currency_balance = self.all.get(base_currency)
+
+        if base_currency_balance is None:
+            total_base_value = portfolio.Amount(base_currency, 0)
+        else:
+            total_base_value = base_currency_balance.exchange_free + \
+                    base_currency_balance.margin_available - \
+                    base_currency_balance.margin_in_position
+
+        amount_in_position = {}
+
+        # Compute balances already in the target position
+        for currency, (ptt, trade_type) in repartition.items():
+            amount_in_position[currency] = portfolio.Amount(base_currency, 0)
+            balance = self.all.get(currency)
+            if currency != base_currency and balance is not None:
+                if trade_type == "short":
+                    amount = balance.margin_borrowed
+                else:
+                    amount = balance.exchange_free + balance.exchange_used
+                amount_in_position[currency] = amount.in_currency(base_currency,
+                    self.market, compute_value=compute_value)
+                total_base_value += amount_in_position[currency]
+
+        # recursively delete more-than-filled positions from the wanted
+        # repartition
+        did_delete = True
+        while did_delete:
+            did_delete = False
+            sum_ratio = sum([v[0] for k, v in repartition.items()])
+            current_base_value = total_base_value
+            for currency, (ptt, trade_type) in repartition.copy().items():
+                if amount_in_position[currency] > current_base_value * ptt / sum_ratio:
+                    did_delete = True
+                    del(repartition[currency])
+                    total_base_value -= amount_in_position[currency]
+        return repartition, total_base_value, amount_in_position
+
     def dispatch_assets(self, amount, liquidity="medium", repartition=None):
         if repartition is None:
             repartition = Portfolio.repartition(liquidity=liquidity)
@@ -521,7 +564,7 @@ class Portfolio:
             cls.retrieve_cryptoportfolio()
         cls.get_cryptoportfolio()
         liquidities = cls.liquidities.get(liquidity)
-        return liquidities[cls.last_date.get()]
+        return liquidities[cls.last_date.get()].copy()
 
     @classmethod
     def get_cryptoportfolio(cls, refetch=False):
diff --git a/tests/test_market.py b/tests/test_market.py
index c89025b..07188ac 100644
--- a/tests/test_market.py
+++ b/tests/test_market.py
@@ -186,14 +186,17 @@ class MarketTest(WebMockTestCase):
                     return { "average": D("0.000001") }
                 if c1 == "ETH" and c2 == "BTC":
                     return { "average": D("0.1") }
+                if c1 == "FOO" and c2 == "BTC":
+                    return { "average": D("0.1") }
                 self.fail("Should not be called with {}, {}".format(c1, c2))
             get_ticker.side_effect = _get_ticker
 
             repartition.return_value = {
-                    "DOGE": (D("0.25"), "short"),
-                    "BTC": (D("0.25"), "long"),
-                    "ETH": (D("0.25"), "long"),
-                    "XMR": (D("0.25"), "long"),
+                    "DOGE": (D("0.20"), "short"),
+                    "BTC": (D("0.20"), "long"),
+                    "ETH": (D("0.20"), "long"),
+                    "XMR": (D("0.20"), "long"),
+                    "FOO": (D("0.20"), "long"),
                     }
             m = market.Market(self.ccxt, self.market_args())
             self.ccxt.fetch_all_balances.return_value = {
@@ -210,12 +213,12 @@ class MarketTest(WebMockTestCase):
                         "total": D("5.0")
                         },
                     "BTC": {
-                        "exchange_free": D("0.075"),
+                        "exchange_free": D("0.065"),
                         "exchange_used": D("0.02"),
-                        "exchange_total": D("0.095"),
-                        "margin_available": D("0.025"),
+                        "exchange_total": D("0.085"),
+                        "margin_available": D("0.035"),
                         "margin_in_position": D("0.01"),
-                        "margin_total": D("0.035"),
+                        "margin_total": D("0.045"),
                         "total": D("0.13")
                         },
                     "ETH": {
@@ -224,6 +227,12 @@ class MarketTest(WebMockTestCase):
                         "exchange_total": D("1.0"),
                         "total": D("1.0")
                         },
+                    "FOO": {
+                        "exchange_free": D("0.1"),
+                        "exchange_used": D("0.0"),
+                        "exchange_total": D("0.1"),
+                        "total": D("0.1"),
+                        },
                     }
 
             m.balances.fetch_balances(tag="tag")
@@ -236,12 +245,13 @@ class MarketTest(WebMockTestCase):
 
             self.assertEqual(portfolio.Amount("BTC", "-0.025"),
                     new_repartition["DOGE"] - values_in_base["DOGE"])
-            self.assertEqual(portfolio.Amount("BTC", "0.025"),
-                    new_repartition["ETH"] - values_in_base["ETH"])
             self.assertEqual(0,
-                    new_repartition["ZRC"] - values_in_base["ZRC"])
+                    new_repartition["ETH"] - values_in_base["ETH"])
+            self.assertIsNone(new_repartition.get("ZRC"))
             self.assertEqual(portfolio.Amount("BTC", "0.025"),
                     new_repartition["XMR"])
+            self.assertEqual(portfolio.Amount("BTC", "0.015"),
+                    new_repartition["FOO"] - values_in_base["FOO"])
 
         compute_trades.reset_mock()
         with self.subTest(available_balance_only=True, balance=0),\
diff --git a/tests/test_store.py b/tests/test_store.py
index d7620a0..1a722b5 100644
--- a/tests/test_store.py
+++ b/tests/test_store.py
@@ -434,6 +434,176 @@ class BalanceStoreTest(WebMockTestCase):
             self.assertListEqual(["XVG", "XMR", "USDT"], list(balance_store.currencies()))
 
     @mock.patch.object(market.Portfolio, "repartition")
+    def test_available_balances_for_repartition(self, repartition):
+        with self.subTest(available_balance_only=True):
+            def _get_ticker(c1, c2):
+                if c1 == "ZRC" and c2 == "BTC":
+                    return { "average": D("0.0001") }
+                if c1 == "DOGE" and c2 == "BTC":
+                    return { "average": D("0.000001") }
+                if c1 == "ETH" and c2 == "BTC":
+                    return { "average": D("0.1") }
+                if c1 == "FOO" and c2 == "BTC":
+                    return { "average": D("0.1") }
+                self.fail("Should not be called with {}, {}".format(c1, c2))
+            self.m.get_ticker.side_effect = _get_ticker
+
+            repartition.return_value = {
+                    "DOGE": (D("0.20"), "short"),
+                    "BTC": (D("0.20"), "long"),
+                    "ETH": (D("0.20"), "long"),
+                    "XMR": (D("0.20"), "long"),
+                    "FOO": (D("0.20"), "long"),
+                    }
+            self.m.ccxt.fetch_all_balances.return_value = {
+                    "ZRC": {
+                        "exchange_free": D("2.0"),
+                        "exchange_used": D("0.0"),
+                        "exchange_total": D("2.0"),
+                        "total": D("2.0")
+                        },
+                    "DOGE": {
+                        "exchange_free": D("5.0"),
+                        "exchange_used": D("0.0"),
+                        "exchange_total": D("5.0"),
+                        "total": D("5.0")
+                        },
+                    "BTC": {
+                        "exchange_free": D("0.065"),
+                        "exchange_used": D("0.02"),
+                        "exchange_total": D("0.085"),
+                        "margin_available": D("0.035"),
+                        "margin_in_position": D("0.01"),
+                        "margin_total": D("0.045"),
+                        "total": D("0.13")
+                        },
+                    "ETH": {
+                        "exchange_free": D("1.0"),
+                        "exchange_used": D("0.0"),
+                        "exchange_total": D("1.0"),
+                        "total": D("1.0")
+                        },
+                    "FOO": {
+                        "exchange_free": D("0.1"),
+                        "exchange_used": D("0.0"),
+                        "exchange_total": D("0.1"),
+                        "total": D("0.1"),
+                        },
+                    }
+
+            balance_store = market.BalanceStore(self.m)
+            balance_store.fetch_balances()
+            _repartition, total_base_value, amount_in_position = balance_store.available_balances_for_repartition()
+            repartition.assert_called_with(liquidity="medium")
+            self.assertEqual((D("0.20"), "short"), _repartition["DOGE"])
+            self.assertEqual((D("0.20"), "long"), _repartition["BTC"])
+            self.assertEqual((D("0.20"), "long"), _repartition["XMR"])
+            self.assertEqual((D("0.20"), "long"), _repartition["FOO"])
+            self.assertIsNone(_repartition.get("ETH"))
+            self.assertEqual(portfolio.Amount("BTC", "0.1"), total_base_value)
+            self.assertEqual(0, amount_in_position["DOGE"])
+            self.assertEqual(0, amount_in_position["BTC"])
+            self.assertEqual(0, amount_in_position["XMR"])
+            self.assertEqual(portfolio.Amount("BTC", "0.1"), amount_in_position["ETH"])
+            self.assertEqual(portfolio.Amount("BTC", "0.01"), amount_in_position["FOO"])
+
+        with self.subTest(available_balance_only=True, balance=0):
+            def _get_ticker(c1, c2):
+                if c1 == "ETH" and c2 == "BTC":
+                    return { "average": D("0.1") }
+                self.fail("Should not be called with {}, {}".format(c1, c2))
+            self.m.get_ticker.side_effect = _get_ticker
+
+            repartition.return_value = {
+                    "BTC": (D("0.5"), "long"),
+                    "ETH": (D("0.5"), "long"),
+                    }
+            self.m.ccxt.fetch_all_balances.return_value = {
+                    "ETH": {
+                        "exchange_free": D("1.0"),
+                        "exchange_used": D("0.0"),
+                        "exchange_total": D("1.0"),
+                        "total": D("1.0")
+                        },
+                    }
+
+            balance_store = market.BalanceStore(self.m)
+            balance_store.fetch_balances()
+            _repartition, total_base_value, amount_in_position = balance_store.available_balances_for_repartition(liquidity="high")
+
+            repartition.assert_called_with(liquidity="high")
+            self.assertEqual((D("0.5"), "long"), _repartition["BTC"])
+            self.assertIsNone(_repartition.get("ETH"))
+            self.assertEqual(0, total_base_value)
+            self.assertEqual(0, amount_in_position["BTC"])
+            self.assertEqual(0, amount_in_position["BTC"])
+
+        repartition.reset_mock()
+        with self.subTest(available_balance_only=True, balance=0,
+                repartition="present"):
+            def _get_ticker(c1, c2):
+                if c1 == "ETH" and c2 == "BTC":
+                    return { "average": D("0.1") }
+                self.fail("Should not be called with {}, {}".format(c1, c2))
+            self.m.get_ticker.side_effect = _get_ticker
+
+            _repartition = {
+                    "BTC": (D("0.5"), "long"),
+                    "ETH": (D("0.5"), "long"),
+                    }
+            self.m.ccxt.fetch_all_balances.return_value = {
+                    "ETH": {
+                        "exchange_free": D("1.0"),
+                        "exchange_used": D("0.0"),
+                        "exchange_total": D("1.0"),
+                        "total": D("1.0")
+                        },
+                    }
+
+            balance_store = market.BalanceStore(self.m)
+            balance_store.fetch_balances()
+            _repartition, total_base_value, amount_in_position = balance_store.available_balances_for_repartition(repartition=_repartition)
+            repartition.assert_not_called()
+
+            self.assertEqual((D("0.5"), "long"), _repartition["BTC"])
+            self.assertIsNone(_repartition.get("ETH"))
+            self.assertEqual(0, total_base_value)
+            self.assertEqual(0, amount_in_position["BTC"])
+            self.assertEqual(portfolio.Amount("BTC", "0.1"), amount_in_position["ETH"])
+
+        repartition.reset_mock()
+        with self.subTest(available_balance_only=True, balance=0,
+                repartition="present", base_currency="ETH"):
+            def _get_ticker(c1, c2):
+                if c1 == "ETH" and c2 == "BTC":
+                    return { "average": D("0.1") }
+                self.fail("Should not be called with {}, {}".format(c1, c2))
+            self.m.get_ticker.side_effect = _get_ticker
+
+            _repartition = {
+                    "BTC": (D("0.5"), "long"),
+                    "ETH": (D("0.5"), "long"),
+                    }
+            self.m.ccxt.fetch_all_balances.return_value = {
+                    "ETH": {
+                        "exchange_free": D("1.0"),
+                        "exchange_used": D("0.0"),
+                        "exchange_total": D("1.0"),
+                        "total": D("1.0")
+                        },
+                    }
+
+            balance_store = market.BalanceStore(self.m)
+            balance_store.fetch_balances()
+            _repartition, total_base_value, amount_in_position = balance_store.available_balances_for_repartition(repartition=_repartition, base_currency="ETH")
+
+            self.assertEqual((D("0.5"), "long"), _repartition["BTC"])
+            self.assertEqual((D("0.5"), "long"), _repartition["ETH"])
+            self.assertEqual(portfolio.Amount("ETH", 1), total_base_value)
+            self.assertEqual(0, amount_in_position["BTC"])
+            self.assertEqual(0, amount_in_position["ETH"])
+
+    @mock.patch.object(market.Portfolio, "repartition")
     def test_dispatch_assets(self, repartition):
         self.m.ccxt.fetch_all_balances.return_value = self.fetch_balance
 
@@ -1343,27 +1513,27 @@ class PortfolioTest(WebMockTestCase):
         with self.subTest(from_cache=False):
             market.Portfolio.liquidities = store.LockedVar({
                     "medium": {
-                        "2018-03-01": "medium_2018-03-01",
-                        "2018-03-08": "medium_2018-03-08",
+                        "2018-03-01": ["medium_2018-03-01"],
+                        "2018-03-08": ["medium_2018-03-08"],
                         },
                     "high": {
-                        "2018-03-01": "high_2018-03-01",
-                        "2018-03-08": "high_2018-03-08",
+                        "2018-03-01": ["high_2018-03-01"],
+                        "2018-03-08": ["high_2018-03-08"],
                         }
                     })
             market.Portfolio.last_date = store.LockedVar("2018-03-08")
 
-            self.assertEqual("medium_2018-03-08", market.Portfolio.repartition())
+            self.assertEqual(["medium_2018-03-08"], market.Portfolio.repartition())
             get_cryptoportfolio.assert_called_once_with()
             retrieve_cryptoportfolio.assert_not_called()
-            self.assertEqual("medium_2018-03-08", market.Portfolio.repartition(liquidity="medium"))
-            self.assertEqual("high_2018-03-08", market.Portfolio.repartition(liquidity="high"))
+            self.assertEqual(["medium_2018-03-08"], market.Portfolio.repartition(liquidity="medium"))
+            self.assertEqual(["high_2018-03-08"], market.Portfolio.repartition(liquidity="high"))
 
         retrieve_cryptoportfolio.reset_mock()
         get_cryptoportfolio.reset_mock()
 
         with self.subTest(from_cache=True):
-            self.assertEqual("medium_2018-03-08", market.Portfolio.repartition(from_cache=True))
+            self.assertEqual(["medium_2018-03-08"], market.Portfolio.repartition(from_cache=True))
             get_cryptoportfolio.assert_called_once_with()
             retrieve_cryptoportfolio.assert_called_once_with()