2 files changed, 135 insertions, 20 deletions

diff --git a/market.py b/market.py
index ce41841..eff670c 100644
--- a/market.py
+++ b/market.py
@@ -210,18 +210,32 @@ class Market:
         self.report.log_stage("follow_orders_end")
 
     def prepare_trades(self, base_currency="BTC", liquidity="medium",
-            compute_value="average", repartition=None, only=None):
+            compute_value="average", repartition=None, only=None,
+            available_balance_only=False):
 
         self.report.log_stage("prepare_trades",
                 base_currency=base_currency, liquidity=liquidity,
                 compute_value=compute_value, only=only,
-                repartition=repartition)
+                repartition=repartition, available_balance_only=available_balance_only)
 
         values_in_base = self.balances.in_currency(base_currency,
                 compute_value=compute_value)
-        total_base_value = sum(values_in_base.values())
+        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
+        else:
+            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
+
         self.trades.compute_trades(values_in_base, new_repartition, only=only)
 
     def print_tickers(self, base_currency="BTC"):
@@ -292,7 +306,7 @@ class Processor:
                     "before": False,
                     "after": True,
                     "fetch_balances": ["begin", "end"],
-                    "prepare_trades": { "only": "acquire" },
+                    "prepare_trades": { "only": "acquire", "available_balance_only": True },
                     "prepare_orders": { "only": "acquire", "compute_value": "average" },
                     "move_balances": {},
                     "run_orders": {},
@@ -326,7 +340,7 @@ class Processor:
                     "before": False,
                     "after": True,
                     "fetch_balances": ["begin", "end"],
-                    "prepare_trades": {},
+                    "prepare_trades": { "available_balance_only": True },
                     "prepare_orders": { "compute_value": "average" },
                     "move_balances": {},
                     "run_orders": {},
diff --git a/tests/test_market.py b/tests/test_market.py
index e0cf70a..53630b7 100644
--- a/tests/test_market.py
+++ b/tests/test_market.py
@@ -130,19 +130,20 @@ class MarketTest(WebMockTestCase):
     @mock.patch.object(market.Market, "get_ticker")
     @mock.patch.object(market.TradeStore, "compute_trades")
     def test_prepare_trades(self, compute_trades, get_ticker, repartition):
-        repartition.return_value = {
-                "XEM": (D("0.75"), "long"),
-                "BTC": (D("0.25"), "long"),
-                }
-        def _get_ticker(c1, c2):
-            if c1 == "USDT" and c2 == "BTC":
-                return { "average": D("0.0001") }
-            if c1 == "XVG" and c2 == "BTC":
-                return { "average": D("0.000001") }
-            self.fail("Should be called with {}, {}".format(c1, c2))
-        get_ticker.side_effect = _get_ticker
-
-        with mock.patch("market.ReportStore"):
+        with self.subTest(available_balance_only=False),\
+                mock.patch("market.ReportStore"):
+            def _get_ticker(c1, c2):
+                if c1 == "USDT" and c2 == "BTC":
+                    return { "average": D("0.0001") }
+                if c1 == "XVG" and c2 == "BTC":
+                    return { "average": D("0.000001") }
+                self.fail("Should not be called with {}, {}".format(c1, c2))
+            get_ticker.side_effect = _get_ticker
+
+            repartition.return_value = {
+                    "XEM": (D("0.75"), "long"),
+                    "BTC": (D("0.25"), "long"),
+                    }
             m = market.Market(self.ccxt, self.market_args())
             self.ccxt.fetch_all_balances.return_value = {
                     "USDT": {
@@ -171,9 +172,109 @@ class MarketTest(WebMockTestCase):
             self.assertEqual(D("0.7575"), call[0][1]["XEM"].value)
             m.report.log_stage.assert_called_once_with("prepare_trades",
                     base_currency='BTC', compute_value='average',
-                    liquidity='medium', only=None, repartition=None)
+                    available_balance_only=False, liquidity='medium',
+                    only=None, repartition=None)
             m.report.log_balances.assert_called_once_with(tag="tag")
 
+        compute_trades.reset_mock()
+        with self.subTest(available_balance_only=True),\
+                mock.patch("market.ReportStore"):
+            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") }
+                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"),
+                    }
+            m = market.Market(self.ccxt, self.market_args())
+            self.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.075"),
+                        "exchange_used": D("0.02"),
+                        "exchange_total": D("0.095"),
+                        "margin_available": D("0.025"),
+                        "margin_in_position": D("0.01"),
+                        "margin_total": D("0.035"),
+                        "total": D("0.13")
+                        },
+                    "ETH": {
+                        "exchange_free": D("1.0"),
+                        "exchange_used": D("0.0"),
+                        "exchange_total": D("1.0"),
+                        "total": D("1.0")
+                        },
+                    }
+
+            m.balances.fetch_balances(tag="tag")
+            m.prepare_trades(available_balance_only=True)
+            compute_trades.assert_called_once()
+
+            call = compute_trades.call_args[0]
+            values_in_base = call[0]
+            new_repartition = call[1]
+
+            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"])
+            self.assertEqual(portfolio.Amount("BTC", "0.025"),
+                    new_repartition["XMR"])
+
+        compute_trades.reset_mock()
+        with self.subTest(available_balance_only=True, balance=0),\
+                mock.patch("market.ReportStore"):
+            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))
+            get_ticker.side_effect = _get_ticker
+
+            repartition.return_value = {
+                    "BTC": (D("0.5"), "long"),
+                    "ETH": (D("0.5"), "long"),
+                    }
+            m = market.Market(self.ccxt, self.market_args())
+            self.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")
+                        },
+                    }
+
+            m.balances.fetch_balances(tag="tag")
+            m.prepare_trades(available_balance_only=True)
+            compute_trades.assert_called_once()
+
+            call = compute_trades.call_args[0]
+            values_in_base = call[0]
+            new_repartition = call[1]
+
+            self.assertEqual(new_repartition["ETH"], values_in_base["ETH"])
 
     @mock.patch.object(market.time, "sleep")
     @mock.patch.object(market.TradeStore, "all_orders")
@@ -859,7 +960,7 @@ class ProcessorTest(WebMockTestCase):
 
         method, arguments = processor.method_arguments("prepare_trades")
         self.assertEqual(m.prepare_trades, method)
-        self.assertEqual(['base_currency', 'liquidity', 'compute_value', 'repartition', 'only'], arguments)
+        self.assertEqual(['base_currency', 'liquidity', 'compute_value', 'repartition', 'only', 'available_balance_only'], arguments)
 
         method, arguments = processor.method_arguments("prepare_orders")
         self.assertEqual(m.trades.prepare_orders, method)