Dynamically use process methods

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):