1 files changed, 44 insertions, 8 deletions

diff --git a/test.py b/test.py
index a4ec8d2..4ed0477 100644
--- a/test.py
+++ b/test.py
@@ -2562,7 +2562,7 @@ class HelperTest(WebMockTestCase):
     @mock.patch("portfolio.Portfolio.wait_for_recent")
     def test_main_process_market(self, wait, buy, sell):
         with self.subTest(before=False, after=False):
-            helper.main_process_market("user")
+            helper.main_process_market("user", None)
             
             wait.assert_not_called()
             buy.assert_not_called()
@@ -2572,7 +2572,7 @@ class HelperTest(WebMockTestCase):
         wait.reset_mock()
         sell.reset_mock()
         with self.subTest(before=True, after=False):
-            helper.main_process_market("user", before=True)
+            helper.main_process_market("user", None, before=True)
             
             wait.assert_not_called()
             buy.assert_not_called()
@@ -2582,7 +2582,7 @@ class HelperTest(WebMockTestCase):
         wait.reset_mock()
         sell.reset_mock()
         with self.subTest(before=False, after=True):
-            helper.main_process_market("user", after=True)
+            helper.main_process_market("user", None, after=True)
             
             wait.assert_called_once_with("user")
             buy.assert_called_once_with("user")
@@ -2592,12 +2592,37 @@ class HelperTest(WebMockTestCase):
         wait.reset_mock()
         sell.reset_mock()
         with self.subTest(before=True, after=True):
-            helper.main_process_market("user", before=True, after=True)
+            helper.main_process_market("user", None, before=True, after=True)
             
             wait.assert_called_once_with("user")
             buy.assert_called_once_with("user")
             sell.assert_called_once_with("user")
 
+        buy.reset_mock()
+        wait.reset_mock()
+        sell.reset_mock()
+        with self.subTest(action="print_balances"),\
+                mock.patch("helper.print_balances") as print_balances:
+            helper.main_process_market("user", "print_balances")
+
+            buy.assert_not_called()
+            wait.assert_not_called()
+            sell.assert_not_called()
+            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")
+
+            buy.assert_not_called()
+            wait.assert_not_called()
+            sell.assert_not_called()
+            print_orders.assert_called_once_with("user")
+
+        with self.subTest(action="unknown"),\
+                self.assertRaises(NotImplementedError):
+            helper.main_process_market("user", "unknown")
+
     @mock.patch.object(helper, "psycopg2")
     def test_fetch_markets(self, psycopg2):
         connect_mock = mock.Mock()
@@ -2607,12 +2632,23 @@ class HelperTest(WebMockTestCase):
         connect_mock.cursor.return_value = cursor_mock
         psycopg2.connect.return_value = connect_mock
 
-        rows = list(helper.main_fetch_markets({"foo": "bar"}))
+        with self.subTest(user=None):
+            rows = list(helper.main_fetch_markets({"foo": "bar"}, None))
+
+            psycopg2.connect.assert_called_once_with(foo="bar")
+            cursor_mock.execute.assert_called_once_with("SELECT config,user_id FROM market_configs")
+
+            self.assertEqual(["row_1", "row_2"], rows)
+
+        psycopg2.connect.reset_mock()
+        cursor_mock.execute.reset_mock()
+        with self.subTest(user=1):
+            rows = list(helper.main_fetch_markets({"foo": "bar"}, 1))
 
-        psycopg2.connect.assert_called_once_with(foo="bar")
-        cursor_mock.execute.assert_called_once_with("SELECT config,user_id FROM market_configs")
+            psycopg2.connect.assert_called_once_with(foo="bar")
+            cursor_mock.execute.assert_called_once_with("SELECT config,user_id FROM market_configs WHERE user_id = %s", 1)
 
-        self.assertEqual(["row_1", "row_2"], rows)
+            self.assertEqual(["row_1", "row_2"], rows)
 
     @mock.patch.object(helper.sys, "exit")
     def test_main_parse_args(self, exit):