From bb127bc87c2b2880469bfab230415c85e589421a Mon Sep 17 00:00:00 2001 From: =?utf8?q?Isma=C3=ABl=20Bouya?= Date: Sat, 5 May 2018 14:41:03 +0200 Subject: [PATCH] Add checkpoints when fetching balance --- market.py | 54 ++++++++++++++++++++++++++++++++------------ store.py | 10 ++++---- tests/test_market.py | 32 ++++++++++++++++++++------ tests/test_store.py | 10 +++++--- 4 files changed, 77 insertions(+), 29 deletions(-) diff --git a/market.py b/market.py index 5876071..7996a58 100644 --- a/market.py +++ b/market.py @@ -254,8 +254,7 @@ class Processor: { "name": "print_balances", "number": 1, - "fetch_balances": ["begin"], - "fetch_balances_args": { "add_portfolio": True }, + "fetch_balances_begin": { "log_tickers": True, "add_portfolio": True }, "print_tickers": { "base_currency": "BTC" }, } ], @@ -272,25 +271,37 @@ class Processor: "number": 2, "before": False, "after": True, - "fetch_balances": ["begin"], + "fetch_balances_begin": {}, "prepare_trades": { "compute_value": "average" }, "prepare_orders": { "compute_value": "average" }, }, ], "sell_needed": [ { - "name": "wait", + "name": "print_balances", "number": 0, + "before": True, + "after": False, + "fetch_balances_begin": { + "checkpoint": "end", + "log_tickers": True, + "add_portfolio": True + }, + }, + { + "name": "wait", + "number": 1, "before": False, "after": True, "wait_for_recent": {}, }, { "name": "sell", - "number": 1, + "number": 2, "before": False, "after": True, - "fetch_balances": ["begin", "end"], + "fetch_balances_begin": {}, + "fetch_balances_end": {}, "prepare_trades": {}, "prepare_orders": { "only": "dispose", "compute_value": "average" }, "run_orders": {}, @@ -299,10 +310,14 @@ class Processor: }, { "name": "buy", - "number": 2, + "number": 3, "before": False, "after": True, - "fetch_balances": ["begin", "end"], + "fetch_balances_begin": {}, + "fetch_balances_end": { + "checkpoint": "begin", + "log_tickers": True + }, "prepare_trades": { "only": "acquire", "available_balance_only": True }, "prepare_orders": { "only": "acquire", "compute_value": "average" }, "move_balances": {}, @@ -317,7 +332,12 @@ class Processor: "number": 1, "before": True, "after": False, - "fetch_balances": ["begin", "end"], + "fetch_balances_begin": { + "checkpoint": "end", + "log_tickers": True, + "add_portfolio": True + }, + "fetch_balances_end": {}, "prepare_trades": { "repartition": { "base_currency": (1, "long") } }, "prepare_orders": { "compute_value": "average" }, "run_orders": {}, @@ -336,7 +356,11 @@ class Processor: "number": 3, "before": False, "after": True, - "fetch_balances": ["begin", "end"], + "fetch_balances_begin": {}, + "fetch_balances_end": { + "checkpoint": "begin", + "log_tickers": True + }, "prepare_trades": { "available_balance_only": True }, "prepare_orders": { "compute_value": "average" }, "move_balances": {}, @@ -388,18 +412,18 @@ class Processor: process_name = "process_{}__{}_{}".format(scenario_name, step["number"], step["name"]) self.market.report.log_stage("{}_begin".format(process_name)) - fetch_args = step.get("fetch_balances_args", {}) - if "begin" in step.get("fetch_balances", []): + if "fetch_balances_begin" in step: self.market.balances.fetch_balances(tag="{}_begin".format(process_name), - log_tickers=True, **fetch_args) + **step["fetch_balances_begin"]) for action in self.ordered_actions: if action in step: self.run_action(action, step[action], kwargs) - if "end" in step.get("fetch_balances", []): + if "fetch_balances_end" in step: self.market.balances.fetch_balances(tag="{}_end".format(process_name), - log_tickers=True, **fetch_args) + **step["fetch_balances_end"]) + self.market.report.log_stage("{}_end".format(process_name)) def method_arguments(self, action): diff --git a/store.py b/store.py index 76cfec8..2a71bed 100644 --- a/store.py +++ b/store.py @@ -99,7 +99,7 @@ class ReportStore: "args": args, }) - def log_balances(self, tag=None, tickers=None, + def log_balances(self, tag=None, checkpoint=None, tickers=None, ticker_currency=None, compute_value=None, type=None): self.print_log("[Balance]") for currency, balance in self.market.balances.all.items(): @@ -108,6 +108,7 @@ class ReportStore: log = { "type": "balance", "tag": tag, + "checkpoint": checkpoint, "balances": self.market.balances.as_json() } @@ -303,7 +304,8 @@ class BalanceStore: compute_value, type) return amounts - def fetch_balances(self, tag=None, add_portfolio=False, log_tickers=False, + def fetch_balances(self, tag=None, add_portfolio=False, + checkpoint=None, log_tickers=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,11 +317,11 @@ class BalanceStore: self.all.setdefault(currency, portfolio.Balance(currency, {})) if log_tickers: tickers = self.in_currency(ticker_currency, compute_value=ticker_compute_value, type=ticker_type) - self.market.report.log_balances(tag=tag, + self.market.report.log_balances(tag=tag, checkpoint=checkpoint, tickers=tickers, ticker_currency=ticker_currency, compute_value=ticker_compute_value, type=ticker_type) else: - self.market.report.log_balances(tag=tag) + self.market.report.log_balances(tag=tag, checkpoint=checkpoint) def dispatch_assets(self, amount, liquidity="medium", repartition=None): if repartition is None: diff --git a/tests/test_market.py b/tests/test_market.py index 46fad53..37c009b 100644 --- a/tests/test_market.py +++ b/tests/test_market.py @@ -174,7 +174,7 @@ class MarketTest(WebMockTestCase): base_currency='BTC', compute_value='average', available_balance_only=False, liquidity='medium', only=None, repartition=None) - m.report.log_balances.assert_called_once_with(tag="tag") + m.report.log_balances.assert_called_once_with(tag="tag", checkpoint=None) compute_trades.reset_mock() with self.subTest(available_balance_only=True),\ @@ -964,7 +964,7 @@ class ProcessorTest(WebMockTestCase): process_step.reset_mock() processor.process("sell_needed", steps=["before", "after"]) - self.assertEqual(3, process_step.call_count) + self.assertEqual(4, process_step.call_count) def test_method_arguments(self): ccxt = mock.Mock(spec=market.ccxt.poloniexE) @@ -1002,17 +1002,17 @@ class ProcessorTest(WebMockTestCase): processor = market.Processor(self.m) with mock.patch.object(processor, "run_action") as run_action: - step = processor.scenarios["sell_needed"][1] + step = processor.scenarios["sell_needed"][2] processor.process_step("foo", step, {"foo":"bar"}) self.m.report.log_stage.assert_has_calls([ - mock.call("process_foo__1_sell_begin"), - mock.call("process_foo__1_sell_end"), + mock.call("process_foo__2_sell_begin"), + mock.call("process_foo__2_sell_end"), ]) self.m.balances.fetch_balances.assert_has_calls([ - mock.call(tag="process_foo__1_sell_begin", log_tickers=True), - mock.call(tag="process_foo__1_sell_end", log_tickers=True), + mock.call(tag="process_foo__2_sell_begin"), + mock.call(tag="process_foo__2_sell_end"), ]) self.assertEqual(5, run_action.call_count) @@ -1029,6 +1029,24 @@ class ProcessorTest(WebMockTestCase): with mock.patch.object(processor, "run_action") as run_action: step = processor.scenarios["sell_needed"][0] + processor.process_step("foo", step, {"foo":"bar"}) + + self.m.report.log_stage.assert_has_calls([ + mock.call("process_foo__0_print_balances_begin"), + mock.call("process_foo__0_print_balances_end"), + ]) + self.m.balances.fetch_balances.assert_has_calls([ + mock.call(add_portfolio=True, checkpoint='end', + log_tickers=True, + tag='process_foo__0_print_balances_begin') + ]) + + self.assertEqual(0, run_action.call_count) + + self.m.reset_mock() + with mock.patch.object(processor, "run_action") as run_action: + step = processor.scenarios["sell_needed"][1] + processor.process_step("foo", step, {"foo":"bar"}) self.m.balances.fetch_balances.assert_not_called() diff --git a/tests/test_store.py b/tests/test_store.py index ee7e063..58e76e0 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -380,7 +380,7 @@ class BalanceStoreTest(WebMockTestCase): balance_store.fetch_balances(tag="foo") self.assertEqual(0, balance_store.all["ETC"].total) self.assertListEqual(["USDT", "XVG", "XMR", "ETC"], list(balance_store.currencies())) - self.m.report.log_balances.assert_called_with(tag="foo") + self.m.report.log_balances.assert_called_with(tag="foo", checkpoint=None) with self.subTest(log_tickers=True),\ mock.patch.object(balance_store, "in_currency") as in_currency: @@ -388,7 +388,7 @@ class BalanceStoreTest(WebMockTestCase): balance_store.fetch_balances(log_tickers=True, ticker_currency="FOO", ticker_compute_value="compute", ticker_type="type") self.m.report.log_balances.assert_called_with(compute_value='compute', - tag=None, ticker_currency='FOO', tickers='tickers', + tag=None, checkpoint=None, ticker_currency='FOO', tickers='tickers', type='type') balance_store = market.BalanceStore(self.m) @@ -425,7 +425,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) - self.m.report.log_balances.assert_called_with(tag=None) + self.m.report.log_balances.assert_called_with(tag=None, checkpoint=None) self.m.report.log_dispatch.assert_called_once_with(portfolio.Amount("BTC", "11.1"), amounts, "medium", repartition_hash) @@ -617,12 +617,14 @@ class ReportStoreTest(WebMockTestCase): ]) add_log.assert_called_once_with({ 'type': 'balance', + 'checkpoint': None, 'balances': 'json', 'tag': 'tag' }) add_redis_status.assert_called_once_with({ 'type': 'balance', 'balances': 'json', + 'checkpoint': None, 'tag': 'tag' }) add_log.reset_mock() @@ -639,6 +641,7 @@ class ReportStoreTest(WebMockTestCase): type="total") add_log.assert_called_once_with({ 'type': 'balance', + 'checkpoint': None, 'balances': 'json', 'tag': 'tag', 'tickers': { @@ -658,6 +661,7 @@ class ReportStoreTest(WebMockTestCase): }) add_redis_status.assert_called_once_with({ 'type': 'balance', + 'checkpoint': None, 'balances': 'json', 'tag': 'tag', 'tickers': { -- 2.41.0