aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--helper.py31
-rw-r--r--main.py4
-rw-r--r--test.py52
3 files changed, 69 insertions, 18 deletions
diff --git a/helper.py b/helper.py
index 4b9ce0d..6d28c3f 100644
--- a/helper.py
+++ b/helper.py
@@ -24,6 +24,11 @@ def main_parse_args(argv):
24 parser.add_argument("--debug", 24 parser.add_argument("--debug",
25 default=False, action='store_const', const=True, 25 default=False, action='store_const', const=True,
26 help="Run in debug mode") 26 help="Run in debug mode")
27 parser.add_argument("--user",
28 default=None, required=False, help="Only run for that user")
29 parser.add_argument("--action",
30 default=None, required=False,
31 help="Do a different action than trading")
27 32
28 args = parser.parse_args(argv) 33 args = parser.parse_args(argv)
29 34
@@ -51,21 +56,31 @@ def main_parse_config(config_file):
51 56
52 return [config["postgresql"], report_path] 57 return [config["postgresql"], report_path]
53 58
54def main_fetch_markets(pg_config): 59def main_fetch_markets(pg_config, user):
55 connection = psycopg2.connect(**pg_config) 60 connection = psycopg2.connect(**pg_config)
56 cursor = connection.cursor() 61 cursor = connection.cursor()
57 62
58 cursor.execute("SELECT config,user_id FROM market_configs") 63 if user is None:
64 cursor.execute("SELECT config,user_id FROM market_configs")
65 else:
66 cursor.execute("SELECT config,user_id FROM market_configs WHERE user_id = %s", user)
59 67
60 for row in cursor: 68 for row in cursor:
61 yield row 69 yield row
62 70
63def main_process_market(user_market, before=False, after=False): 71def main_process_market(user_market, action, before=False, after=False):
64 if before: 72 if action is None:
65 process_sell_all__1_all_sell(user_market) 73 if before:
66 if after: 74 process_sell_all__1_all_sell(user_market)
67 portfolio.Portfolio.wait_for_recent(user_market) 75 if after:
68 process_sell_all__2_all_buy(user_market) 76 portfolio.Portfolio.wait_for_recent(user_market)
77 process_sell_all__2_all_buy(user_market)
78 elif action == "print_balances":
79 print_balances(user_market)
80 elif action == "print_orders":
81 print_orders(user_market)
82 else:
83 raise NotImplementedError("Unknown action {}".format(action))
69 84
70def main_store_report(report_path, user_id, user_market): 85def main_store_report(report_path, user_id, user_market):
71 try: 86 try:
diff --git a/main.py b/main.py
index e7cdcf0..3cb7f4a 100644
--- a/main.py
+++ b/main.py
@@ -5,11 +5,11 @@ args = helper.main_parse_args(sys.argv[1:])
5 5
6pg_config, report_path = helper.main_parse_config(args.config) 6pg_config, report_path = helper.main_parse_config(args.config)
7 7
8for market_config, user_id in helper.main_fetch_markets(pg_config): 8for market_config, user_id in helper.main_fetch_markets(pg_config, args.user):
9 try: 9 try:
10 market_config["apiKey"] = market_config.pop("key") 10 market_config["apiKey"] = market_config.pop("key")
11 user_market = market.Market.from_config(market_config, debug=args.debug) 11 user_market = market.Market.from_config(market_config, debug=args.debug)
12 helper.main_process_market(user_market, before=args.before, after=args.after) 12 helper.main_process_market(user_market, args.action, before=args.before, after=args.after)
13 except Exception as e: 13 except Exception as e:
14 print("{}: {}".format(e.__class__.__name__, e)) 14 print("{}: {}".format(e.__class__.__name__, e))
15 finally: 15 finally:
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):
2562 @mock.patch("portfolio.Portfolio.wait_for_recent") 2562 @mock.patch("portfolio.Portfolio.wait_for_recent")
2563 def test_main_process_market(self, wait, buy, sell): 2563 def test_main_process_market(self, wait, buy, sell):
2564 with self.subTest(before=False, after=False): 2564 with self.subTest(before=False, after=False):
2565 helper.main_process_market("user") 2565 helper.main_process_market("user", None)
2566 2566
2567 wait.assert_not_called() 2567 wait.assert_not_called()
2568 buy.assert_not_called() 2568 buy.assert_not_called()
@@ -2572,7 +2572,7 @@ class HelperTest(WebMockTestCase):
2572 wait.reset_mock() 2572 wait.reset_mock()
2573 sell.reset_mock() 2573 sell.reset_mock()
2574 with self.subTest(before=True, after=False): 2574 with self.subTest(before=True, after=False):
2575 helper.main_process_market("user", before=True) 2575 helper.main_process_market("user", None, before=True)
2576 2576
2577 wait.assert_not_called() 2577 wait.assert_not_called()
2578 buy.assert_not_called() 2578 buy.assert_not_called()
@@ -2582,7 +2582,7 @@ class HelperTest(WebMockTestCase):
2582 wait.reset_mock() 2582 wait.reset_mock()
2583 sell.reset_mock() 2583 sell.reset_mock()
2584 with self.subTest(before=False, after=True): 2584 with self.subTest(before=False, after=True):
2585 helper.main_process_market("user", after=True) 2585 helper.main_process_market("user", None, after=True)
2586 2586
2587 wait.assert_called_once_with("user") 2587 wait.assert_called_once_with("user")
2588 buy.assert_called_once_with("user") 2588 buy.assert_called_once_with("user")
@@ -2592,12 +2592,37 @@ class HelperTest(WebMockTestCase):
2592 wait.reset_mock() 2592 wait.reset_mock()
2593 sell.reset_mock() 2593 sell.reset_mock()
2594 with self.subTest(before=True, after=True): 2594 with self.subTest(before=True, after=True):
2595 helper.main_process_market("user", before=True, after=True) 2595 helper.main_process_market("user", None, before=True, after=True)
2596 2596
2597 wait.assert_called_once_with("user") 2597 wait.assert_called_once_with("user")
2598 buy.assert_called_once_with("user") 2598 buy.assert_called_once_with("user")
2599 sell.assert_called_once_with("user") 2599 sell.assert_called_once_with("user")
2600 2600
2601 buy.reset_mock()
2602 wait.reset_mock()
2603 sell.reset_mock()
2604 with self.subTest(action="print_balances"),\
2605 mock.patch("helper.print_balances") as print_balances:
2606 helper.main_process_market("user", "print_balances")
2607
2608 buy.assert_not_called()
2609 wait.assert_not_called()
2610 sell.assert_not_called()
2611 print_balances.assert_called_once_with("user")
2612
2613 with self.subTest(action="print_orders"),\
2614 mock.patch("helper.print_orders") as print_orders:
2615 helper.main_process_market("user", "print_orders")
2616
2617 buy.assert_not_called()
2618 wait.assert_not_called()
2619 sell.assert_not_called()
2620 print_orders.assert_called_once_with("user")
2621
2622 with self.subTest(action="unknown"),\
2623 self.assertRaises(NotImplementedError):
2624 helper.main_process_market("user", "unknown")
2625
2601 @mock.patch.object(helper, "psycopg2") 2626 @mock.patch.object(helper, "psycopg2")
2602 def test_fetch_markets(self, psycopg2): 2627 def test_fetch_markets(self, psycopg2):
2603 connect_mock = mock.Mock() 2628 connect_mock = mock.Mock()
@@ -2607,12 +2632,23 @@ class HelperTest(WebMockTestCase):
2607 connect_mock.cursor.return_value = cursor_mock 2632 connect_mock.cursor.return_value = cursor_mock
2608 psycopg2.connect.return_value = connect_mock 2633 psycopg2.connect.return_value = connect_mock
2609 2634
2610 rows = list(helper.main_fetch_markets({"foo": "bar"})) 2635 with self.subTest(user=None):
2636 rows = list(helper.main_fetch_markets({"foo": "bar"}, None))
2637
2638 psycopg2.connect.assert_called_once_with(foo="bar")
2639 cursor_mock.execute.assert_called_once_with("SELECT config,user_id FROM market_configs")
2640
2641 self.assertEqual(["row_1", "row_2"], rows)
2642
2643 psycopg2.connect.reset_mock()
2644 cursor_mock.execute.reset_mock()
2645 with self.subTest(user=1):
2646 rows = list(helper.main_fetch_markets({"foo": "bar"}, 1))
2611 2647
2612 psycopg2.connect.assert_called_once_with(foo="bar") 2648 psycopg2.connect.assert_called_once_with(foo="bar")
2613 cursor_mock.execute.assert_called_once_with("SELECT config,user_id FROM market_configs") 2649 cursor_mock.execute.assert_called_once_with("SELECT config,user_id FROM market_configs WHERE user_id = %s", 1)
2614 2650
2615 self.assertEqual(["row_1", "row_2"], rows) 2651 self.assertEqual(["row_1", "row_2"], rows)
2616 2652
2617 @mock.patch.object(helper.sys, "exit") 2653 @mock.patch.object(helper.sys, "exit")
2618 def test_main_parse_args(self, exit): 2654 def test_main_parse_args(self, exit):