Add USDT rate to balances

4 files changed, 45 insertions, 4 deletions

diff --git a/market.py b/market.py
index b521ea9..caa9513 100644
--- a/market.py
+++ b/market.py
@@ -257,7 +257,11 @@ class Processor:
                 {
                     "name": "print_balances",
                     "number": 1,
-                    "fetch_balances_begin": { "log_tickers": True, "add_portfolio": True },
+                    "fetch_balances_begin": {
+                        "log_tickers": True,
+                        "add_usdt": True,
+                        "add_portfolio": True
+                        },
                     "print_tickers": { "base_currency": "BTC" },
                     }
                 ],
@@ -288,6 +292,7 @@ class Processor:
                     "fetch_balances_begin": {
                         "checkpoint": "end",
                         "log_tickers": True,
+                        "add_usdt": True,
                         "add_portfolio": True
                         },
                     },
@@ -319,6 +324,7 @@ class Processor:
                     "fetch_balances_begin": {},
                     "fetch_balances_end": {
                         "checkpoint": "begin",
+                        "add_usdt": True,
                         "log_tickers": True
                         },
                     "prepare_trades": { "only": "acquire", "available_balance_only": True },
@@ -338,6 +344,7 @@ class Processor:
                     "fetch_balances_begin": {
                         "checkpoint": "end",
                         "log_tickers": True,
+                        "add_usdt": True,
                         "add_portfolio": True
                         },
                     "fetch_balances_end": {},
@@ -362,6 +369,7 @@ class Processor:
                     "fetch_balances_begin": {},
                     "fetch_balances_end": {
                         "checkpoint": "begin",
+                        "add_usdt": True,
                         "log_tickers": True
                         },
                     "prepare_trades": { "available_balance_only": True },
diff --git a/store.py b/store.py
index 2a71bed..32c4121 100644
--- a/store.py
+++ b/store.py
@@ -305,7 +305,7 @@ class BalanceStore:
         return amounts
 
     def fetch_balances(self, tag=None, add_portfolio=False,
-            checkpoint=None, log_tickers=False,
+            checkpoint=None, log_tickers=False, add_usdt=False,
             ticker_currency="BTC", ticker_compute_value="average", ticker_type="total"):
         all_balances = self.market.ccxt.fetch_all_balances()
         for currency, balance in all_balances.items():
@@ -315,6 +315,8 @@ class BalanceStore:
         if add_portfolio:
             for currency in Portfolio.repartition(from_cache=True):
                 self.all.setdefault(currency, portfolio.Balance(currency, {}))
+        if add_usdt:
+            self.all.setdefault("USDT", portfolio.Balance("USDT", {}))
         if log_tickers:
             tickers = self.in_currency(ticker_currency, compute_value=ticker_compute_value, type=ticker_type)
             self.market.report.log_balances(tag=tag, checkpoint=checkpoint,
diff --git a/tests/test_market.py b/tests/test_market.py
index 6a9ea76..2c92655 100644
--- a/tests/test_market.py
+++ b/tests/test_market.py
@@ -1042,6 +1042,7 @@ class ProcessorTest(WebMockTestCase):
             self.m.balances.fetch_balances.assert_has_calls([
                 mock.call(add_portfolio=True, checkpoint='end',
                     log_tickers=True,
+                    add_usdt=True,
                     tag='process_foo__0_print_balances_begin')
                 ])
 
@@ -1060,7 +1061,7 @@ class ProcessorTest(WebMockTestCase):
 
             processor.process_step("foo", step, {"foo":"bar"})
             self.m.balances.fetch_balances.assert_called_once_with(
-                    add_portfolio=True, log_tickers=True,
+                    add_portfolio=True, add_usdt=True, log_tickers=True,
                     tag='process_foo__1_print_balances_begin')
 
     def test_parse_args(self):
diff --git a/tests/test_store.py b/tests/test_store.py
index 58e76e0..d7620a0 100644
--- a/tests/test_store.py
+++ b/tests/test_store.py
@@ -392,7 +392,6 @@ class BalanceStoreTest(WebMockTestCase):
                     type='type')
 
         balance_store = market.BalanceStore(self.m)
-
         with self.subTest(add_portfolio=True),\
                 mock.patch.object(market.Portfolio, "repartition") as repartition:
             repartition.return_value = {
@@ -402,6 +401,37 @@ class BalanceStoreTest(WebMockTestCase):
             balance_store.fetch_balances(add_portfolio=True)
             self.assertListEqual(["USDT", "XVG", "XMR", "DOGE"], list(balance_store.currencies()))
 
+        self.m.ccxt.fetch_all_balances.return_value = {
+                "ETC": {
+                    "exchange_free": 0,
+                    "exchange_used": 0,
+                    "exchange_total": 0,
+                    "margin_total": 0,
+                    },
+                "XVG": {
+                    "exchange_free": 16,
+                    "exchange_used": 0,
+                    "exchange_total": 16,
+                    "margin_total": 0,
+                    },
+                "XMR": {
+                    "exchange_free": 0,
+                    "exchange_used": 0,
+                    "exchange_total": 0,
+                    "margin_total": D("-1.0"),
+                    "margin_free": 0,
+                    },
+                }
+
+        balance_store = market.BalanceStore(self.m)
+        with self.subTest(add_usdt=True),\
+                mock.patch.object(market.Portfolio, "repartition") as repartition:
+            repartition.return_value = {
+                    "DOGE": D("0.5"),
+                    "ETH": D("0.5"),
+                    }
+            balance_store.fetch_balances(add_usdt=True)
+            self.assertListEqual(["XVG", "XMR", "USDT"], list(balance_store.currencies()))
 
     @mock.patch.object(market.Portfolio, "repartition")
     def test_dispatch_assets(self, repartition):