Fix available balance when buying

2 files changed, 199 insertions, 19 deletions

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()