From 23f4616a85bef7e9acc57740f889a2a0346788af Mon Sep 17 00:00:00 2001 From: =?utf8?q?Isma=C3=ABl=20Bouya?= Date: Thu, 1 Mar 2018 15:23:27 +0100 Subject: [PATCH] Dynamically use process methods --- helper.py | 18 +++++++++--------- test.py | 10 ++++++---- 2 files changed, 15 insertions(+), 13 deletions(-) diff --git a/helper.py b/helper.py index d948dac..21e95c7 100644 --- a/helper.py +++ b/helper.py @@ -83,8 +83,8 @@ def main_parse_args(argv): parser.add_argument("--user", default=None, required=False, help="Only run for that user") parser.add_argument("--action", - default=None, required=False, - help="Do a different action than trading") + action='append', + help="Do a different action than trading (add several times to chain)") args = parser.parse_args(argv) @@ -124,19 +124,19 @@ def main_fetch_markets(pg_config, user): for row in cursor: yield row -def main_process_market(user_market, action, before=False, after=False): - if action is None: +def main_process_market(user_market, actions, before=False, after=False): + if len(actions or []) == 0: if before: process_sell_all__1_all_sell(user_market) if after: portfolio.Portfolio.wait_for_recent(user_market) process_sell_all__2_all_buy(user_market) - elif action == "print_balances": - print_balances(user_market) - elif action == "print_orders": - print_orders(user_market) else: - raise NotImplementedError("Unknown action {}".format(action)) + for action in actions: + if action in globals(): + (globals()[action])(user_market) + else: + raise NotImplementedError("Unknown action {}".format(action)) def main_store_report(report_path, user_id, user_market): try: diff --git a/test.py b/test.py index fbc857d..14b5559 100644 --- a/test.py +++ b/test.py @@ -2920,7 +2920,7 @@ class HelperTest(WebMockTestCase): sell.reset_mock() with self.subTest(action="print_balances"),\ mock.patch("helper.print_balances") as print_balances: - helper.main_process_market("user", "print_balances") + helper.main_process_market("user", ["print_balances"]) buy.assert_not_called() wait.assert_not_called() @@ -2928,17 +2928,19 @@ class HelperTest(WebMockTestCase): print_balances.assert_called_once_with("user") with self.subTest(action="print_orders"),\ - mock.patch("helper.print_orders") as print_orders: - helper.main_process_market("user", "print_orders") + mock.patch("helper.print_orders") as print_orders,\ + mock.patch("helper.print_balances") as print_balances: + helper.main_process_market("user", ["print_orders", "print_balances"]) buy.assert_not_called() wait.assert_not_called() sell.assert_not_called() print_orders.assert_called_once_with("user") + print_balances.assert_called_once_with("user") with self.subTest(action="unknown"),\ self.assertRaises(NotImplementedError): - helper.main_process_market("user", "unknown") + helper.main_process_market("user", ["unknown"]) @mock.patch.object(helper, "psycopg2") def test_fetch_markets(self, psycopg2): -- 2.41.0