1 files changed, 54 insertions, 19 deletions

diff --git a/test.py b/test.py
index 0f1f130..34b2fe3 100644
--- a/test.py
+++ b/test.py
@@ -653,6 +653,8 @@ class HelperTest(WebMockTestCase):
                     "total": D("10000.0")
                     },
                 }
+        portfolio.BalanceStore.fetch_balances(market, tag="tag")
+
         helper.prepare_trades(market)
         compute_trades.assert_called()
 
@@ -663,7 +665,7 @@ class HelperTest(WebMockTestCase):
         self.assertEqual(D("0.2525"), call[0][1]["BTC"].value)
         self.assertEqual(D("0.7575"), call[0][1]["XEM"].value)
         report_store_h.log_stage.assert_called_once_with("prepare_trades")
-        report_store.log_balances.assert_called_once_with(market)
+        report_store.log_balances.assert_called_once_with(market, tag="tag")
 
     @mock.patch.object(portfolio.Portfolio, "repartition")
     @mock.patch.object(helper, "get_ticker")
@@ -700,6 +702,8 @@ class HelperTest(WebMockTestCase):
                     "total": D("10000.0")
                     },
                 }
+        portfolio.BalanceStore.fetch_balances(market, tag="tag")
+
         helper.update_trades(market)
         compute_trades.assert_called()
 
@@ -710,7 +714,7 @@ class HelperTest(WebMockTestCase):
         self.assertEqual(D("0.2525"), call[0][1]["BTC"].value)
         self.assertEqual(D("0.7575"), call[0][1]["XEM"].value)
         report_store_h.log_stage.assert_called_once_with("update_trades")
-        report_store.log_balances.assert_called_once_with(market)
+        report_store.log_balances.assert_called_once_with(market, tag="tag")
 
     @mock.patch.object(portfolio.Portfolio, "repartition")
     @mock.patch.object(helper, "get_ticker")
@@ -741,6 +745,8 @@ class HelperTest(WebMockTestCase):
                     "total": D("10000.0")
                     },
                 }
+        portfolio.BalanceStore.fetch_balances(market, tag="tag")
+
         helper.prepare_trades_to_sell_all(market)
         repartition.assert_not_called()
         compute_trades.assert_called()
@@ -751,7 +757,7 @@ class HelperTest(WebMockTestCase):
         self.assertEqual(D("0.01"), call[0][0]["XVG"].value)
         self.assertEqual(D("1.01"), call[0][1]["BTC"].value)
         report_store_h.log_stage.assert_called_once_with("prepare_trades_to_sell_all")
-        report_store.log_balances.assert_called_once_with(market)
+        report_store.log_balances.assert_called_once_with(market, tag="tag")
 
     @mock.patch.object(portfolio.time, "sleep")
     @mock.patch.object(portfolio.TradeStore, "all_orders")
@@ -875,8 +881,9 @@ class HelperTest(WebMockTestCase):
 
     @mock.patch.object(helper, "prepare_trades")
     @mock.patch.object(portfolio.TradeStore, "prepare_orders")
+    @mock.patch.object(portfolio.BalanceStore, "fetch_balances")
     @mock.patch.object(portfolio.ReportStore, "log_stage")
-    def test_print_orders(self, log_stage, prepare_orders, prepare_trades):
+    def test_print_orders(self, log_stage, fetch_balances, prepare_orders, prepare_trades):
         market = mock.Mock()
         portfolio.BalanceStore.all = {
                 "BTC": portfolio.Balance("BTC", {
@@ -890,7 +897,9 @@ class HelperTest(WebMockTestCase):
                     "exchange_free": 3,
                     "exchange_used": 0}),
                 }
+
         helper.print_orders(market)
+        fetch_balances.assert_called_with(market, tag="print_orders")
         prepare_trades.assert_called_with(market, base_currency="BTC",
                 compute_value="average", debug=True)
         prepare_orders.assert_called_with(compute_value="average")
@@ -928,9 +937,11 @@ class HelperTest(WebMockTestCase):
     @mock.patch.object(helper, "follow_orders")
     @mock.patch.object(portfolio.TradeStore, "prepare_orders")
     @mock.patch.object(portfolio.TradeStore, "run_orders")
+    @mock.patch.object(portfolio.BalanceStore, "fetch_balances")
     @mock.patch.object(portfolio.ReportStore, "log_stage")
-    def test_process_sell_needed__1_sell(self, log_stage, run_orders,
-            prepare_orders, follow_orders, prepare_trades):
+    def test_process_sell_needed__1_sell(self, log_stage,
+            fetch_balances, run_orders, prepare_orders, follow_orders,
+            prepare_trades):
         market = mock.Mock()
         portfolio.BalanceStore.all = {
                 "BTC": portfolio.Balance("BTC", {
@@ -945,6 +956,10 @@ class HelperTest(WebMockTestCase):
                     "exchange_used": 0}),
                 }
         helper.process_sell_needed__1_sell(market)
+        fetch_balances.assert_has_calls([
+            mock.call(market, tag="process_sell_needed__1_sell_begin"),
+            mock.call(market, tag="process_sell_needed__1_sell_end"),
+            ])
         prepare_trades.assert_called_with(market, base_currency="BTC",
                 liquidity="medium", debug=False)
         prepare_orders.assert_called_with(compute_value="average",
@@ -958,10 +973,11 @@ class HelperTest(WebMockTestCase):
     @mock.patch.object(helper, "move_balances")
     @mock.patch.object(portfolio.TradeStore, "prepare_orders")
     @mock.patch.object(portfolio.TradeStore, "run_orders")
+    @mock.patch.object(portfolio.BalanceStore, "fetch_balances")
     @mock.patch.object(portfolio.ReportStore, "log_stage")
-    def test_process_sell_needed__2_buy(self, log_stage, run_orders,
-            prepare_orders, move_balances,
-            follow_orders, update_trades):
+    def test_process_sell_needed__2_buy(self, log_stage, fetch_balances,
+            run_orders, prepare_orders, move_balances, follow_orders,
+            update_trades):
         market = mock.Mock()
         portfolio.BalanceStore.all = {
                 "BTC": portfolio.Balance("BTC", {
@@ -976,6 +992,10 @@ class HelperTest(WebMockTestCase):
                     "exchange_used": 0}),
                 }
         helper.process_sell_needed__2_buy(market)
+        fetch_balances.assert_has_calls([
+            mock.call(market, tag="process_sell_needed__2_buy_begin"),
+            mock.call(market, tag="process_sell_needed__2_buy_end"),
+            ])
         update_trades.assert_called_with(market, base_currency="BTC",
                 debug=False, liquidity="medium", only="acquire")
         prepare_orders.assert_called_with(compute_value="average",
@@ -989,9 +1009,11 @@ class HelperTest(WebMockTestCase):
     @mock.patch.object(helper, "follow_orders")
     @mock.patch.object(portfolio.TradeStore, "prepare_orders")
     @mock.patch.object(portfolio.TradeStore, "run_orders")
+    @mock.patch.object(portfolio.BalanceStore, "fetch_balances")
     @mock.patch.object(portfolio.ReportStore, "log_stage")
-    def test_process_sell_all__1_sell(self, log_stage, run_orders,
-            prepare_orders, follow_orders, prepare_trades_to_sell_all):
+    def test_process_sell_all__1_sell(self, log_stage, fetch_balances,
+            run_orders, prepare_orders, follow_orders,
+            prepare_trades_to_sell_all):
         market = mock.Mock()
         portfolio.BalanceStore.all = {
                 "BTC": portfolio.Balance("BTC", {
@@ -1006,6 +1028,10 @@ class HelperTest(WebMockTestCase):
                     "exchange_used": 0}),
                 }
         helper.process_sell_all__1_all_sell(market)
+        fetch_balances.assert_has_calls([
+            mock.call(market, tag="process_sell_all__1_all_sell_begin"),
+            mock.call(market, tag="process_sell_all__1_all_sell_end"),
+            ])
         prepare_trades_to_sell_all.assert_called_with(market, base_currency="BTC",
                 debug=False)
         prepare_orders.assert_called_with(compute_value="average")
@@ -1018,10 +1044,11 @@ class HelperTest(WebMockTestCase):
     @mock.patch.object(helper, "move_balances")
     @mock.patch.object(portfolio.TradeStore, "prepare_orders")
     @mock.patch.object(portfolio.TradeStore, "run_orders")
+    @mock.patch.object(portfolio.BalanceStore, "fetch_balances")
     @mock.patch.object(portfolio.ReportStore, "log_stage")
-    def test_process_sell_all__2_all_buy(self, log_stage, run_orders,
-            prepare_orders, move_balances, follow_orders,
-            prepare_trades):
+    def test_process_sell_all__2_all_buy(self, log_stage,
+            fetch_balances, run_orders, prepare_orders, move_balances,
+            follow_orders, prepare_trades):
         market = mock.Mock()
         portfolio.BalanceStore.all = {
                 "BTC": portfolio.Balance("BTC", {
@@ -1036,6 +1063,10 @@ class HelperTest(WebMockTestCase):
                     "exchange_used": 0}),
                 }
         helper.process_sell_all__2_all_buy(market)
+        fetch_balances.assert_has_calls([
+            mock.call(market, tag="process_sell_all__2_all_buy_begin"),
+            mock.call(market, tag="process_sell_all__2_all_buy_end"),
+            ])
         prepare_trades.assert_called_with(market, base_currency="BTC",
                 liquidity="medium", debug=False)
         prepare_orders.assert_called_with(compute_value="average")
@@ -1308,10 +1339,10 @@ class BalanceStoreTest(WebMockTestCase):
         portfolio.BalanceStore.all["ETC"] = portfolio.Balance("ETC", {
             "exchange_total": "1", "exchange_free": "0",
             "exchange_used": "1" })
-        portfolio.BalanceStore.fetch_balances(market)
+        portfolio.BalanceStore.fetch_balances(market, tag="foo")
         self.assertEqual(0, portfolio.BalanceStore.all["ETC"].total)
         self.assertListEqual(["USDT", "XVG", "XMR", "ETC"], list(portfolio.BalanceStore.currencies()))
-        log_balances.assert_called_with(market)
+        log_balances.assert_called_with(market, tag="foo")
 
     @mock.patch.object(portfolio.Portfolio, "repartition")
     @mock.patch.object(portfolio.ReportStore, "log_balances")
@@ -1336,7 +1367,7 @@ class BalanceStoreTest(WebMockTestCase):
         self.assertEqual(D("2.6"), amounts["BTC"].value)
         self.assertEqual(D("7.5"), amounts["XEM"].value)
         self.assertEqual(D("-1.0"), amounts["DASH"].value)
-        log_balances.assert_called_with(market)
+        log_balances.assert_called_with(market, tag=None)
         log_dispatch.assert_called_once_with(portfolio.Amount("BTC",
             "11.1"), amounts, "medium", repartition_hash)
 
@@ -2317,13 +2348,17 @@ class ReportStoreTest(WebMockTestCase):
         balance_store.as_json.return_value = "json"
         balance_store.all = { "FOO": "bar", "BAR": "baz" }
 
-        portfolio.ReportStore.log_balances("market")
+        portfolio.ReportStore.log_balances("market", tag="tag")
         print_log.assert_has_calls([
             mock.call("[Balance]"),
             mock.call("\tbar"),
             mock.call("\tbaz"),
             ])
-        add_log.assert_called_once_with({'type': 'balance', 'balances': 'json'})
+        add_log.assert_called_once_with({
+            'type': 'balance',
+            'balances': 'json',
+            'tag': 'tag'
+            })
 
     @mock.patch.object(portfolio.ReportStore, "print_log")
     @mock.patch.object(portfolio.ReportStore, "add_log")