Handle timeouts for move_balances

2 files changed, 149 insertions, 6 deletions

diff --git a/market.py b/market.py
index 055967c..ca65bca 100644
--- a/market.py
+++ b/market.py
@@ -1,10 +1,11 @@
-from ccxt import ExchangeError, NotSupported
+from ccxt import ExchangeError, NotSupported, RequestTimeout
 import ccxt_wrapper as ccxt
 import time
 import psycopg2
 from store import *
 from cachetools.func import ttl_cache
 from datetime import datetime
+from retry import retry
 import portfolio
 
 class Market:
@@ -88,6 +89,7 @@ class Market:
         finally:
             self.store_report()
 
+    @retry(RequestTimeout, tries=5)
     def move_balances(self):
         needed_in_margin = {} 
         moving_to_margin = {}
@@ -102,13 +104,21 @@ class Market:
             current_balance = self.balances.all[currency].margin_available
             moving_to_margin[currency] = (needed - current_balance)
             delta = moving_to_margin[currency].value
+            action = "Moving {} from exchange to margin".format(moving_to_margin[currency])
+
             if self.debug and delta != 0:
-                self.report.log_debug_action("Moving {} from exchange to margin".format(moving_to_margin[currency]))
+                self.report.log_debug_action(action)
                 continue
-            if delta > 0:
-                self.ccxt.transfer_balance(currency, delta, "exchange", "margin")
-            elif delta < 0:
-                self.ccxt.transfer_balance(currency, -delta, "margin", "exchange")
+            try:
+                if delta > 0:
+                    self.ccxt.transfer_balance(currency, delta, "exchange", "margin")
+                elif delta < 0:
+                    self.ccxt.transfer_balance(currency, -delta, "margin", "exchange")
+            except RequestTimeout as e:
+                self.report.log_error(action, message="Retrying", exception=e)
+                self.report.log_move_balances(needed_in_margin, moving_to_margin)
+                self.balances.fetch_balances()
+                raise e
         self.report.log_move_balances(needed_in_margin, moving_to_margin)
 
         self.balances.fetch_balances()
diff --git a/test.py b/test.py
index 18616c1..ea1fd9a 100644
--- a/test.py
+++ b/test.py
@@ -1444,6 +1444,139 @@ class MarketTest(WebMockTestCase):
                     self.ccxt.transfer_balance.assert_any_call("USDT", 100, "exchange", "margin")
                     self.ccxt.transfer_balance.assert_any_call("ETC", 5, "margin", "exchange")
 
+        m.report.reset_mock()
+        fetch_balances.reset_mock()
+        with self.subTest(retry=True):
+            with mock.patch("market.ReportStore"):
+                m = market.Market(self.ccxt, self.market_args())
+
+                value_from = portfolio.Amount("BTC", "0.0")
+                value_from.linked_to = portfolio.Amount("ETH", "0.0")
+                value_to = portfolio.Amount("BTC", "-3.0")
+                trade = portfolio.Trade(value_from, value_to, "ETH", m)
+
+                m.trades.all = [trade]
+                balance = portfolio.Balance("BTC", { "margin_in_position": "0", "margin_available": "0" })
+                m.balances.all = {"BTC": balance}
+
+                m.ccxt.transfer_balance.side_effect = [
+                        market.ccxt.RequestTimeout,
+                        True
+                        ]
+                m.move_balances()
+                self.ccxt.transfer_balance.assert_has_calls([
+                    mock.call("BTC", 3, "exchange", "margin"),
+                    mock.call("BTC", 3, "exchange", "margin")
+                    ])
+                self.assertEqual(2, fetch_balances.call_count)
+                m.report.log_error.assert_called_with(mock.ANY, message="Retrying", exception=mock.ANY)
+                self.assertEqual(2, m.report.log_move_balances.call_count)
+
+        self.ccxt.transfer_balance.reset_mock()
+        m.report.reset_mock()
+        fetch_balances.reset_mock()
+        with self.subTest(retry=True, too_much=True):
+            with mock.patch("market.ReportStore"):
+                m = market.Market(self.ccxt, self.market_args())
+
+                value_from = portfolio.Amount("BTC", "0.0")
+                value_from.linked_to = portfolio.Amount("ETH", "0.0")
+                value_to = portfolio.Amount("BTC", "-3.0")
+                trade = portfolio.Trade(value_from, value_to, "ETH", m)
+
+                m.trades.all = [trade]
+                balance = portfolio.Balance("BTC", { "margin_in_position": "0", "margin_available": "0" })
+                m.balances.all = {"BTC": balance}
+
+                m.ccxt.transfer_balance.side_effect = [
+                        market.ccxt.RequestTimeout,
+                        market.ccxt.RequestTimeout,
+                        market.ccxt.RequestTimeout,
+                        market.ccxt.RequestTimeout,
+                        market.ccxt.RequestTimeout,
+                        ]
+                with self.assertRaises(market.ccxt.RequestTimeout):
+                    m.move_balances()
+
+        self.ccxt.transfer_balance.reset_mock()
+        m.report.reset_mock()
+        fetch_balances.reset_mock()
+        with self.subTest(retry=True, partial_result=True):
+            with mock.patch("market.ReportStore"):
+                m = market.Market(self.ccxt, self.market_args())
+
+                value_from = portfolio.Amount("BTC", "1.0")
+                value_from.linked_to = portfolio.Amount("ETH", "10.0")
+                value_to = portfolio.Amount("BTC", "10.0")
+                trade1 = portfolio.Trade(value_from, value_to, "ETH", m)
+
+                value_from = portfolio.Amount("BTC", "0.0")
+                value_from.linked_to = portfolio.Amount("ETH", "0.0")
+                value_to = portfolio.Amount("BTC", "-3.0")
+                trade2 = portfolio.Trade(value_from, value_to, "ETH", m)
+
+                value_from = portfolio.Amount("USDT", "0.0")
+                value_from.linked_to = portfolio.Amount("XVG", "0.0")
+                value_to = portfolio.Amount("USDT", "-50.0")
+                trade3 = portfolio.Trade(value_from, value_to, "XVG", m)
+
+                m.trades.all = [trade1, trade2, trade3]
+                balance1 = portfolio.Balance("BTC", { "margin_in_position": "0", "margin_available": "0" })
+                balance2 = portfolio.Balance("USDT", { "margin_in_position": "100", "margin_available": "50" })
+                balance3 = portfolio.Balance("ETC", { "margin_in_position": "10", "margin_available": "15" })
+                m.balances.all = {"BTC": balance1, "USDT": balance2, "ETC": balance3}
+
+                call_counts = { "BTC": 0, "USDT": 0, "ETC": 0 }
+                def _transfer_balance(currency, amount, from_, to_):
+                    call_counts[currency] += 1
+                    if currency == "BTC":
+                        m.balances.all["BTC"] = portfolio.Balance("BTC", { "margin_in_position": "0", "margin_available": "3" })
+                    if currency == "USDT":
+                        if call_counts["USDT"] == 1:
+                            raise market.ccxt.RequestTimeout
+                        else:
+                            m.balances.all["USDT"] = portfolio.Balance("USDT", { "margin_in_position": "100", "margin_available": "150" })
+                    if currency == "ETC":
+                            m.balances.all["ETC"] = portfolio.Balance("ETC", { "margin_in_position": "10", "margin_available": "10" })
+
+
+                m.ccxt.transfer_balance.side_effect = _transfer_balance
+
+                m.move_balances()
+                self.ccxt.transfer_balance.assert_has_calls([
+                    mock.call("BTC", 3, "exchange", "margin"),
+                    mock.call('USDT', 100, 'exchange', 'margin'),
+                    mock.call('USDT', 100, 'exchange', 'margin'),
+                    mock.call("ETC", 5, "margin", "exchange")
+                    ])
+                self.assertEqual(2, fetch_balances.call_count)
+                m.report.log_error.assert_called_with(mock.ANY, message="Retrying", exception=mock.ANY)
+                self.assertEqual(2, m.report.log_move_balances.call_count)
+                m.report.log_move_balances.asser_has_calls([
+                    mock.call(
+                        {
+                            'BTC': portfolio.Amount("BTC", "3"),
+                            'USDT': portfolio.Amount("USDT", "150"),
+                            'ETC': portfolio.Amount("ETC", "10"),
+                            },
+                        {
+                            'BTC': portfolio.Amount("BTC", "3"),
+                            'USDT': portfolio.Amount("USDT", "100"),
+                            }),
+                    mock.call(
+                        {
+                            'BTC': portfolio.Amount("BTC", "3"),
+                            'USDT': portfolio.Amount("USDT", "150"),
+                            'ETC': portfolio.Amount("ETC", "10"),
+                            },
+                        {
+                            'BTC': portfolio.Amount("BTC", "0"),
+                            'USDT': portfolio.Amount("USDT", "100"),
+                            'ETC': portfolio.Amount("ETC", "-5"),
+                            }),
+                    ])
+
+
     def test_store_file_report(self):
         file_open = mock.mock_open()
         m = market.Market(self.ccxt, self.market_args(), report_path="present", user_id=1)