Handle timeouts for move_balances

[perso/Immae/Projets/Cryptomonnaies/Cryptoportfolio/Trader.git] / test.py

diff --git a/test.py b/test.py
index 18616c1c848620d93a19d61c72b4edcd7c773ea1..ea1fd9a72c17da48fd398d68504125e5b04bbed7 100644 (file)
--- 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)