Add tags to balance log

3 files changed, 67 insertions, 25 deletions

diff --git a/helper.py b/helper.py
index 87539f1..d9c69cc 100644
--- a/helper.py
+++ b/helper.py
@@ -74,7 +74,6 @@ def fetch_fees(market):
 
 def prepare_trades(market, base_currency="BTC", liquidity="medium", compute_value="average", debug=False):
     ReportStore.log_stage("prepare_trades")
-    BalanceStore.fetch_balances(market)
     values_in_base = BalanceStore.in_currency(base_currency, market, compute_value=compute_value)
     total_base_value = sum(values_in_base.values())
     new_repartition = BalanceStore.dispatch_assets(total_base_value, liquidity=liquidity)
@@ -84,7 +83,6 @@ def prepare_trades(market, base_currency="BTC", liquidity="medium", compute_valu
 
 def update_trades(market, base_currency="BTC", liquidity="medium", compute_value="average", only=None, debug=False):
     ReportStore.log_stage("update_trades")
-    BalanceStore.fetch_balances(market)
     values_in_base = BalanceStore.in_currency(base_currency, market, compute_value=compute_value)
     total_base_value = sum(values_in_base.values())
     new_repartition = BalanceStore.dispatch_assets(total_base_value, liquidity=liquidity)
@@ -92,7 +90,6 @@ def update_trades(market, base_currency="BTC", liquidity="medium", compute_value
 
 def prepare_trades_to_sell_all(market, base_currency="BTC", compute_value="average", debug=False):
     ReportStore.log_stage("prepare_trades_to_sell_all")
-    BalanceStore.fetch_balances(market)
     values_in_base = BalanceStore.in_currency(base_currency, market, compute_value=compute_value)
     total_base_value = sum(values_in_base.values())
     new_repartition = BalanceStore.dispatch_assets(total_base_value, repartition={ base_currency: (1, "long") })
@@ -120,6 +117,7 @@ def follow_orders(sleep=None):
 
 def print_orders(market, base_currency="BTC"):
     ReportStore.log_stage("print_orders")
+    BalanceStore.fetch_balances(market, tag="print_orders")
     prepare_trades(market, base_currency=base_currency, compute_value="average", debug=True)
     TradeStore.prepare_orders(compute_value="average")
 
@@ -131,36 +129,44 @@ def print_balances(market, base_currency="BTC"):
 
 def process_sell_needed__1_sell(market, liquidity="medium", base_currency="BTC", debug=False):
     ReportStore.log_stage("process_sell_needed__1_sell_begin")
+    BalanceStore.fetch_balances(market, tag="process_sell_needed__1_sell_begin")
     prepare_trades(market, liquidity=liquidity, base_currency=base_currency, debug=debug)
     TradeStore.prepare_orders(compute_value="average", only="dispose")
     TradeStore.run_orders()
     follow_orders()
+    BalanceStore.fetch_balances(market, tag="process_sell_needed__1_sell_end")
     ReportStore.log_stage("process_sell_needed__1_sell_end")
 
 def process_sell_needed__2_buy(market, liquidity="medium", base_currency="BTC", debug=False):
     ReportStore.log_stage("process_sell_needed__2_buy_begin")
+    BalanceStore.fetch_balances(market, tag="process_sell_needed__2_buy_begin")
     update_trades(market, base_currency=base_currency, liquidity=liquidity, debug=debug, only="acquire")
     TradeStore.prepare_orders(compute_value="average", only="acquire")
     move_balances(market, debug=debug)
     TradeStore.run_orders()
     follow_orders()
+    BalanceStore.fetch_balances(market, tag="process_sell_needed__2_buy_end")
     ReportStore.log_stage("process_sell_needed__2_buy_end")
 
 def process_sell_all__1_all_sell(market, base_currency="BTC", debug=False, liquidity="medium"):
     ReportStore.log_stage("process_sell_all__1_all_sell_begin")
+    BalanceStore.fetch_balances(market, tag="process_sell_all__1_all_sell_begin")
     prepare_trades_to_sell_all(market, base_currency=base_currency, debug=debug)
     TradeStore.prepare_orders(compute_value="average")
     TradeStore.run_orders()
     follow_orders()
+    BalanceStore.fetch_balances(market, tag="process_sell_all__1_all_sell_end")
     ReportStore.log_stage("process_sell_all__1_all_sell_end")
 
 def process_sell_all__2_all_buy(market, base_currency="BTC", debug=False, liquidity="medium"):
     ReportStore.log_stage("process_sell_all__2_all_buy_begin")
+    BalanceStore.fetch_balances(market, tag="process_sell_all__2_all_buy_begin")
     prepare_trades(market, liquidity=liquidity, base_currency=base_currency, debug=debug)
     TradeStore.prepare_orders(compute_value="average")
     move_balances(market, debug=debug)
     TradeStore.run_orders()
     follow_orders()
+    BalanceStore.fetch_balances(market, tag="process_sell_all__2_all_buy_end")
     ReportStore.log_stage("process_sell_all__2_all_buy_end")
 
 
diff --git a/store.py b/store.py
index dfadef2..a7aad22 100644
--- a/store.py
+++ b/store.py
@@ -43,13 +43,14 @@ class ReportStore:
             })
 
     @classmethod
-    def log_balances(cls, market):
+    def log_balances(cls, market, tag=None):
         cls.print_log("[Balance]")
         for currency, balance in BalanceStore.all.items():
             cls.print_log("\t{}".format(balance))
 
         cls.add_log({
             "type": "balance",
+            "tag": tag,
             "balances": BalanceStore.as_json()
             })
 
@@ -196,13 +197,13 @@ class BalanceStore:
         return amounts
 
     @classmethod
-    def fetch_balances(cls, market):
+    def fetch_balances(cls, market, tag=None):
         all_balances = market.fetch_all_balances()
         for currency, balance in all_balances.items():
             if balance["exchange_total"] != 0 or balance["margin_total"] != 0 or \
                     currency in cls.all:
                 cls.all[currency] = portfolio.Balance(currency, balance)
-        ReportStore.log_balances(market)
+        ReportStore.log_balances(market, tag=tag)
 
     @classmethod
     def dispatch_assets(cls, amount, liquidity="medium", repartition=None):
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")