Use generic parser for market processing

3 files changed, 76 insertions, 93 deletions

diff --git a/main.py b/main.py
index 2c7b570..f465d8d 100644
--- a/main.py
+++ b/main.py
@@ -149,7 +149,10 @@ def parse_args(argv):
     parser.add_argument("--db-database", default="cryptoportfolio",
             help="Database access to database (default: cryptoportfolio)")
 
-    return parser.parse_args(argv)
+    parsed = parser.parse_args(argv)
+    if parsed.action is None:
+        parsed.action = ["sell_all"]
+    return parsed
 
 def process(market_config, market_id, user_id, args, pg_config):
     try:
diff --git a/market.py b/market.py
index 7a37cf6..4593eb1 100644
--- a/market.py
+++ b/market.py
@@ -76,17 +76,13 @@ class Market:
 
     def process(self, actions, before=False, after=False):
         try:
-            if len(actions or []) == 0:
-                if before:
-                    self.processor.process("sell_all", steps="before")
-                if after:
-                    self.processor.process("sell_all", steps="after")
-            else:
-                for action in actions:
-                    if hasattr(self, action):
-                        getattr(self, action)()
-                    else:
-                        self.report.log_error("market_process", message="Unknown action {}".format(action))
+            for action in actions:
+                if bool(before) is bool(after):
+                    self.processor.process(action, steps="all")
+                elif before:
+                    self.processor.process(action, steps="before")
+                elif after:
+                    self.processor.process(action, steps="after")
         except Exception as e:
             self.report.log_error("market_process", exception=e)
         finally:
@@ -212,16 +208,7 @@ class Market:
                 liquidity=liquidity, repartition=repartition)
         self.trades.compute_trades(values_in_base, new_repartition, only=only)
 
-    # Helpers
-    def print_orders(self, base_currency="BTC"):
-        self.report.log_stage("print_orders")
-        self.balances.fetch_balances(tag="print_orders")
-        self.prepare_trades(base_currency=base_currency, compute_value="average")
-        self.trades.prepare_orders(compute_value="average")
-
-    def print_balances(self, base_currency="BTC"):
-        self.report.log_stage("print_balances")
-        self.balances.fetch_balances()
+    def print_tickers(self, base_currency="BTC"):
         if base_currency is not None:
             self.report.print_log("total:")
             self.report.print_log(sum(self.balances.in_currency(base_currency).values()))
@@ -237,12 +224,20 @@ class Processor:
                     "wait_for_recent": {},
                     },
                 ],
+            "print_balances": [
+                {
+                    "name": "print_balances",
+                    "number": 1,
+                    "fetch_balances": ["begin"],
+                    "print_tickers": { "base_currency": "BTC" },
+                    }
+                ],
             "print_orders": [
                 {
                     "name": "wait",
                     "number": 1,
-                    "before": False,
-                    "after": True,
+                    "before": True,
+                    "after": False,
                     "wait_for_recent": {},
                     },
                 {
@@ -328,7 +323,7 @@ class Processor:
     ordered_actions = [
             "wait_for_recent", "prepare_trades", "prepare_orders",
             "move_balances", "run_orders", "follow_orders",
-            "close_trades"]
+            "close_trades", "print_tickers"]
 
     def __init__(self, market):
         self.market = market
@@ -337,7 +332,7 @@ class Processor:
         if step == "all":
             return scenario
         elif step == "before" or step == "after":
-            return list(filter(lambda x: step in x and x[step], scenario))
+            return list(filter(lambda x: x.get(step, False), scenario))
         elif type(step) == int:
             return [scenario[step-1]]
         elif type(step) == str:
@@ -345,7 +340,12 @@ class Processor:
         else:
             raise TypeError("Unknown step {}".format(step))
 
+    def can_process(self, scenario_name):
+        return scenario_name in self.scenarios
+
     def process(self, scenario_name, steps="all", **kwargs):
+        if not self.can_process(scenario_name):
+            raise TypeError("Unknown scenario {}".format(scenario_name))
         scenario = self.scenarios[scenario_name]
         selected_steps = []
 
@@ -388,6 +388,8 @@ class Processor:
             method = self.market.follow_orders
         elif action == "close_trades":
             method = self.market.trades.close_trades
+        elif action == "print_tickers":
+            method = self.market.print_tickers
 
         signature = inspect.getfullargspec(method)
         defaults = signature.defaults or []
diff --git a/tests/test_market.py b/tests/test_market.py
index 14b23b5..b41cd6a 100644
--- a/tests/test_market.py
+++ b/tests/test_market.py
@@ -607,21 +607,7 @@ class MarketTest(WebMockTestCase):
             file_report.assert_called_once_with(datetime.datetime(2018, 2, 25))
             db_report.assert_called_once_with(datetime.datetime(2018, 2, 25))
 
-    def test_print_orders(self):
-        m = market.Market(self.ccxt, self.market_args())
-        with mock.patch.object(m.report, "log_stage") as log_stage,\
-                mock.patch.object(m.balances, "fetch_balances") as fetch_balances,\
-                mock.patch.object(m, "prepare_trades") as prepare_trades,\
-                mock.patch.object(m.trades, "prepare_orders") as prepare_orders:
-            m.print_orders()
-
-            log_stage.assert_called_with("print_orders")
-            fetch_balances.assert_called_with(tag="print_orders")
-            prepare_trades.assert_called_with(base_currency="BTC",
-                    compute_value="average")
-            prepare_orders.assert_called_with(compute_value="average")
-
-    def test_print_balances(self):
+    def test_print_tickers(self):
         m = market.Market(self.ccxt, self.market_args())
 
         with mock.patch.object(m.balances, "in_currency") as in_currency,\
@@ -634,10 +620,8 @@ class MarketTest(WebMockTestCase):
                     "ETH": portfolio.Amount("BTC", "0.3"),
                     }
 
-            m.print_balances()
+            m.print_tickers()
 
-            log_stage.assert_called_once_with("print_balances")
-            fetch_balances.assert_called_with()
             print_log.assert_has_calls([
                 mock.call("total:"),
                 mock.call(portfolio.Amount("BTC", "0.95")),
@@ -648,8 +632,8 @@ class MarketTest(WebMockTestCase):
     @mock.patch("market.Market.store_report")
     def test_process(self, store_report, log_error, process):
         m = market.Market(self.ccxt, self.market_args())
-        with self.subTest(before=False, after=False):
-            m.process(None)
+        with self.subTest(actions=[], before=False, after=False):
+            m.process([])
 
             process.assert_not_called()
             store_report.assert_called_once()
@@ -659,9 +643,9 @@ class MarketTest(WebMockTestCase):
         log_error.reset_mock()
         store_report.reset_mock()
         with self.subTest(before=True, after=False):
-            m.process(None, before=True)
+            m.process(["foo"], before=True)
 
-            process.assert_called_once_with("sell_all", steps="before")
+            process.assert_called_once_with("foo", steps="before")
             store_report.assert_called_once()
             log_error.assert_not_called()
 
@@ -669,7 +653,7 @@ class MarketTest(WebMockTestCase):
         log_error.reset_mock()
         store_report.reset_mock()
         with self.subTest(before=False, after=True):
-            m.process(None, after=True)
+            m.process(["sell_all"], after=True)
 
             process.assert_called_once_with("sell_all", steps="after")
             store_report.assert_called_once()
@@ -678,54 +662,30 @@ class MarketTest(WebMockTestCase):
         process.reset_mock()
         log_error.reset_mock()
         store_report.reset_mock()
-        with self.subTest(before=True, after=True):
-            m.process(None, before=True, after=True)
+        with self.subTest(before=False, after=False):
+            m.process(["foo"])
 
-            process.assert_has_calls([
-                mock.call("sell_all", steps="before"),
-                mock.call("sell_all", steps="after"),
-                ])
+            process.assert_called_once_with("foo", steps="all")
             store_report.assert_called_once()
             log_error.assert_not_called()
 
         process.reset_mock()
         log_error.reset_mock()
         store_report.reset_mock()
-        with self.subTest(action="print_balances"),\
-                mock.patch.object(m, "print_balances") as print_balances:
-            m.process(["print_balances"])
+        with self.subTest(before=True, after=True):
+            m.process(["sell_all"], before=True, after=True)
 
-            process.assert_not_called()
-            log_error.assert_not_called()
+            process.assert_called_once_with("sell_all", steps="all")
             store_report.assert_called_once()
-            print_balances.assert_called_once_with()
-
-        log_error.reset_mock()
-        store_report.reset_mock()
-        with self.subTest(action="print_orders"),\
-                mock.patch.object(m, "print_orders") as print_orders,\
-                mock.patch.object(m, "print_balances") as print_balances:
-            m.process(["print_orders", "print_balances"])
-
-            process.assert_not_called()
             log_error.assert_not_called()
-            store_report.assert_called_once()
-            print_orders.assert_called_once_with()
-            print_balances.assert_called_once_with()
-
-        log_error.reset_mock()
-        store_report.reset_mock()
-        with self.subTest(action="unknown"):
-            m.process(["unknown"])
-            log_error.assert_called_once_with("market_process", message="Unknown action unknown")
-            store_report.assert_called_once()
 
+        process.reset_mock()
         log_error.reset_mock()
         store_report.reset_mock()
         with self.subTest(unhandled_exception=True):
             process.side_effect = Exception("bouh")
 
-            m.process(None, before=True)
+            m.process(["some_action"], before=True)
             log_error.assert_called_with("market_process", exception=mock.ANY)
             store_report.assert_called_once()
  
@@ -768,24 +728,39 @@ class ProcessorTest(WebMockTestCase):
         with self.assertRaises(TypeError):
             processor.select_steps(scenario, ["wait"])
 
+    def test_can_process(self):
+        processor = market.Processor(self.m)
+
+        with self.subTest(True):
+            self.assertTrue(processor.can_process("sell_all"))
+
+        with self.subTest(False):
+            self.assertFalse(processor.can_process("unknown_action"))
+
     @mock.patch("market.Processor.process_step")
     def test_process(self, process_step):
-        processor = market.Processor(self.m)
+        with self.subTest("unknown action"):
+            processor = market.Processor(self.m)
+            with self.assertRaises(TypeError):
+                processor.process("unknown_action")
+
+        with self.subTest("nominal case"):
+            processor = market.Processor(self.m)
 
-        processor.process("sell_all", foo="bar")
-        self.assertEqual(3, process_step.call_count)
+            processor.process("sell_all", foo="bar")
+            self.assertEqual(3, process_step.call_count)
 
-        steps = list(map(lambda x: x[1][1]["name"], process_step.mock_calls))
-        scenario_names = list(map(lambda x: x[1][0], process_step.mock_calls))
-        kwargs = list(map(lambda x: x[1][2], process_step.mock_calls))
-        self.assertEqual(["all_sell", "wait", "all_buy"], steps)
-        self.assertEqual(["sell_all", "sell_all", "sell_all"], scenario_names)
-        self.assertEqual([{"foo":"bar"}, {"foo":"bar"}, {"foo":"bar"}], kwargs)
+            steps = list(map(lambda x: x[1][1]["name"], process_step.mock_calls))
+            scenario_names = list(map(lambda x: x[1][0], process_step.mock_calls))
+            kwargs = list(map(lambda x: x[1][2], process_step.mock_calls))
+            self.assertEqual(["all_sell", "wait", "all_buy"], steps)
+            self.assertEqual(["sell_all", "sell_all", "sell_all"], scenario_names)
+            self.assertEqual([{"foo":"bar"}, {"foo":"bar"}, {"foo":"bar"}], kwargs)
 
-        process_step.reset_mock()
+            process_step.reset_mock()
 
-        processor.process("sell_needed", steps=["before", "after"])
-        self.assertEqual(3, process_step.call_count)
+            processor.process("sell_needed", steps=["before", "after"])
+            self.assertEqual(3, process_step.call_count)
 
     def test_method_arguments(self):
         ccxt = mock.Mock(spec=market.ccxt.poloniexE)
@@ -816,6 +791,9 @@ class ProcessorTest(WebMockTestCase):
         method, arguments = processor.method_arguments("close_trades")
         self.assertEqual(m.trades.close_trades, method)
 
+        method, arguments = processor.method_arguments("print_tickers")
+        self.assertEqual(m.print_tickers, method)
+
     def test_process_step(self):
         processor = market.Processor(self.m)