1 files changed, 63 insertions, 36 deletions

diff --git a/tests/test_market.py b/tests/test_market.py
index 6a3322c..46fad53 100644
--- a/tests/test_market.py
+++ b/tests/test_market.py
@@ -1,5 +1,5 @@
 from .helper import *
-import market, store, portfolio
+import market, store, portfolio, dbs
 import datetime
 
 @unittest.skipUnless("unit" in limits, "Unit skipped")
@@ -595,13 +595,11 @@ class MarketTest(WebMockTestCase):
 
             self.assertRegex(stdout_mock.getvalue(), "impossible to store report file: FileNotFoundError;")
 
-    @mock.patch.object(market, "psycopg2")
-    def test_store_database_report(self, psycopg2):
-        connect_mock = mock.Mock()
+    @mock.patch.object(dbs, "psql")
+    def test_store_database_report(self, psql):
         cursor_mock = mock.MagicMock()
 
-        connect_mock.cursor.return_value = cursor_mock
-        psycopg2.connect.return_value = connect_mock
+        psql.cursor.return_value = cursor_mock
         m = market.Market(self.ccxt, self.market_args(),
                 pg_config={"config": "pg_config"}, user_id=1)
         cursor_mock.fetchone.return_value = [42]
@@ -613,7 +611,7 @@ class MarketTest(WebMockTestCase):
                     ("date2", "type2", "payload2"),
                     ]
             m.store_database_report(datetime.datetime(2018, 3, 24))
-            connect_mock.assert_has_calls([
+            psql.assert_has_calls([
                 mock.call.cursor(),
                 mock.call.cursor().execute('INSERT INTO reports("date", "market_config_id", "debug") VALUES (%s, %s, %s) RETURNING id;', (datetime.datetime(2018, 3, 24), None, False)),
                 mock.call.cursor().fetchone(),
@@ -621,21 +619,16 @@ class MarketTest(WebMockTestCase):
                 mock.call.cursor().execute('INSERT INTO report_lines("date", "report_id", "type", "payload") VALUES (%s, %s, %s, %s);', ('date2', 42, 'type2', 'payload2')),
                 mock.call.commit(),
                 mock.call.cursor().close(),
-                mock.call.close()
                 ])
 
-        connect_mock.reset_mock()
         with self.subTest(error=True),\
                 mock.patch('sys.stdout', new_callable=StringIO) as stdout_mock:
-            psycopg2.connect.side_effect = Exception("Bouh")
+            psql.cursor.side_effect = Exception("Bouh")
             m.store_database_report(datetime.datetime(2018, 3, 24))
             self.assertEqual(stdout_mock.getvalue(), "impossible to store report to database: Exception; Bouh\n")
 
-    @mock.patch.object(market, "redis")
+    @mock.patch.object(dbs, "redis")
     def test_store_redis_report(self, redis):
-        connect_mock = mock.Mock()
-        redis.Redis.return_value = connect_mock
-
         m = market.Market(self.ccxt, self.market_args(),
                 redis_config={"config": "redis_config"}, market_id=1)
 
@@ -646,7 +639,7 @@ class MarketTest(WebMockTestCase):
                     ("type2", "payload2"),
                     ]
             m.store_redis_report(datetime.datetime(2018, 3, 24))
-            connect_mock.assert_has_calls([
+            redis.assert_has_calls([
                 mock.call.set("/cryptoportfolio/1/2018-03-24T00:00:00/type1", "payload1", ex=31*24*60*60),
                 mock.call.set("/cryptoportfolio/1/latest/type1", "payload1"),
                 mock.call.set("/cryptoportfolio/1/2018-03-24T00:00:00/type2", "payload2", ex=31*24*60*60),
@@ -654,20 +647,24 @@ class MarketTest(WebMockTestCase):
                 mock.call.set("/cryptoportfolio/1/latest/date", "2018-03-24T00:00:00"),
                 ])
 
-        connect_mock.reset_mock()
+        redis.reset_mock()
         with self.subTest(error=True),\
                 mock.patch('sys.stdout', new_callable=StringIO) as stdout_mock:
-            redis.Redis.side_effect = Exception("Bouh")
+            redis.set.side_effect = Exception("Bouh")
             m.store_redis_report(datetime.datetime(2018, 3, 24))
             self.assertEqual(stdout_mock.getvalue(), "impossible to store report to redis: Exception; Bouh\n")
 
     def test_store_report(self):
         m = market.Market(self.ccxt, self.market_args(report_db=False), user_id=1)
-        with self.subTest(file=None, pg_config=None),\
+        with self.subTest(file=None, pg_connected=None),\
+                mock.patch.object(dbs, "psql_connected") as psql,\
+                mock.patch.object(dbs, "redis_connected") as redis,\
                 mock.patch.object(m, "report") as report,\
                 mock.patch.object(m, "store_database_report") as db_report,\
                 mock.patch.object(m, "store_redis_report") as redis_report,\
                 mock.patch.object(m, "store_file_report") as file_report:
+            psql.return_value = False
+            redis.return_value = False
             m.store_report()
             report.merge.assert_called_with(store.Portfolio.report)
 
@@ -677,13 +674,16 @@ class MarketTest(WebMockTestCase):
 
         report.reset_mock()
         m = market.Market(self.ccxt, self.market_args(report_db=False, report_path="present"), user_id=1)
-        with self.subTest(file="present", pg_config=None),\
+        with self.subTest(file="present", pg_connected=None),\
+                mock.patch.object(dbs, "psql_connected") as psql,\
+                mock.patch.object(dbs, "redis_connected") as redis,\
                 mock.patch.object(m, "report") as report,\
                 mock.patch.object(m, "store_file_report") as file_report,\
                 mock.patch.object(m, "store_redis_report") as redis_report,\
                 mock.patch.object(m, "store_database_report") as db_report,\
                 mock.patch.object(market.datetime, "datetime") as time_mock:
-
+            psql.return_value = False
+            redis.return_value = False
             time_mock.now.return_value = datetime.datetime(2018, 2, 25)
 
             m.store_report()
@@ -695,13 +695,16 @@ class MarketTest(WebMockTestCase):
 
         report.reset_mock()
         m = market.Market(self.ccxt, self.market_args(report_db=True, report_path="present"), user_id=1)
-        with self.subTest(file="present", pg_config=None, report_db=True),\
+        with self.subTest(file="present", pg_connected=None, report_db=True),\
+                mock.patch.object(dbs, "psql_connected") as psql,\
+                mock.patch.object(dbs, "redis_connected") as redis,\
                 mock.patch.object(m, "report") as report,\
                 mock.patch.object(m, "store_file_report") as file_report,\
                 mock.patch.object(m, "store_redis_report") as redis_report,\
                 mock.patch.object(m, "store_database_report") as db_report,\
                 mock.patch.object(market.datetime, "datetime") as time_mock:
-
+            psql.return_value = False
+            redis.return_value = False
             time_mock.now.return_value = datetime.datetime(2018, 2, 25)
 
             m.store_report()
@@ -712,14 +715,17 @@ class MarketTest(WebMockTestCase):
             redis_report.assert_not_called()
 
         report.reset_mock()
-        m = market.Market(self.ccxt, self.market_args(report_db=True), pg_config="present", user_id=1)
-        with self.subTest(file=None, pg_config="present"),\
+        m = market.Market(self.ccxt, self.market_args(report_db=True), user_id=1)
+        with self.subTest(file=None, pg_connected=True),\
+                mock.patch.object(dbs, "psql_connected") as psql,\
+                mock.patch.object(dbs, "redis_connected") as redis,\
                 mock.patch.object(m, "report") as report,\
                 mock.patch.object(m, "store_file_report") as file_report,\
                 mock.patch.object(m, "store_redis_report") as redis_report,\
                 mock.patch.object(m, "store_database_report") as db_report,\
                 mock.patch.object(market.datetime, "datetime") as time_mock:
-
+            psql.return_value = True
+            redis.return_value = False
             time_mock.now.return_value = datetime.datetime(2018, 2, 25)
 
             m.store_report()
@@ -731,14 +737,17 @@ class MarketTest(WebMockTestCase):
 
         report.reset_mock()
         m = market.Market(self.ccxt, self.market_args(report_db=True, report_path="present"),
-                pg_config="pg_config", user_id=1)
-        with self.subTest(file="present", pg_config="present"),\
+                user_id=1)
+        with self.subTest(file="present", pg_connected=True),\
+                mock.patch.object(dbs, "psql_connected") as psql,\
+                mock.patch.object(dbs, "redis_connected") as redis,\
                 mock.patch.object(m, "report") as report,\
                 mock.patch.object(m, "store_file_report") as file_report,\
                 mock.patch.object(m, "store_redis_report") as redis_report,\
                 mock.patch.object(m, "store_database_report") as db_report,\
                 mock.patch.object(market.datetime, "datetime") as time_mock:
-
+            psql.return_value = True
+            redis.return_value = False
             time_mock.now.return_value = datetime.datetime(2018, 2, 25)
 
             m.store_report()
@@ -750,14 +759,17 @@ class MarketTest(WebMockTestCase):
 
         report.reset_mock()
         m = market.Market(self.ccxt, self.market_args(report_redis=False),
-                redis_config="redis_config", user_id=1)
-        with self.subTest(redis_config="present", report_redis=False),\
+                user_id=1)
+        with self.subTest(redis_connected=True, report_redis=False),\
+                mock.patch.object(dbs, "psql_connected") as psql,\
+                mock.patch.object(dbs, "redis_connected") as redis,\
                 mock.patch.object(m, "report") as report,\
                 mock.patch.object(m, "store_file_report") as file_report,\
                 mock.patch.object(m, "store_redis_report") as redis_report,\
                 mock.patch.object(m, "store_database_report") as db_report,\
                 mock.patch.object(market.datetime, "datetime") as time_mock:
-
+            psql.return_value = False
+            redis.return_value = True
             time_mock.now.return_value = datetime.datetime(2018, 2, 25)
 
             m.store_report()
@@ -766,13 +778,16 @@ class MarketTest(WebMockTestCase):
         report.reset_mock()
         m = market.Market(self.ccxt, self.market_args(report_redis=True),
                 user_id=1)
-        with self.subTest(redis_config="absent", report_redis=True),\
+        with self.subTest(redis_connected=False, report_redis=True),\
+                mock.patch.object(dbs, "psql_connected") as psql,\
+                mock.patch.object(dbs, "redis_connected") as redis,\
                 mock.patch.object(m, "report") as report,\
                 mock.patch.object(m, "store_file_report") as file_report,\
                 mock.patch.object(m, "store_redis_report") as redis_report,\
                 mock.patch.object(m, "store_database_report") as db_report,\
                 mock.patch.object(market.datetime, "datetime") as time_mock:
-
+            psql.return_value = False
+            redis.return_value = False
             time_mock.now.return_value = datetime.datetime(2018, 2, 25)
 
             m.store_report()
@@ -780,14 +795,17 @@ class MarketTest(WebMockTestCase):
 
         report.reset_mock()
         m = market.Market(self.ccxt, self.market_args(report_redis=True),
-                redis_config="redis_config", user_id=1)
-        with self.subTest(redis_config="present", report_redis=True),\
+                user_id=1)
+        with self.subTest(redis_connected=True, report_redis=True),\
+                mock.patch.object(dbs, "psql_connected") as psql,\
+                mock.patch.object(dbs, "redis_connected") as redis,\
                 mock.patch.object(m, "report") as report,\
                 mock.patch.object(m, "store_file_report") as file_report,\
                 mock.patch.object(m, "store_redis_report") as redis_report,\
                 mock.patch.object(m, "store_database_report") as db_report,\
                 mock.patch.object(market.datetime, "datetime") as time_mock:
-
+            psql.return_value = False
+            redis.return_value = True
             time_mock.now.return_value = datetime.datetime(2018, 2, 25)
 
             m.store_report()
@@ -1014,6 +1032,15 @@ class ProcessorTest(WebMockTestCase):
             processor.process_step("foo", step, {"foo":"bar"})
             self.m.balances.fetch_balances.assert_not_called()
 
+        self.m.reset_mock()
+        with mock.patch.object(processor, "run_action") as run_action:
+            step = processor.scenarios["print_balances"][0]
+
+            processor.process_step("foo", step, {"foo":"bar"})
+            self.m.balances.fetch_balances.assert_called_once_with(
+                    add_portfolio=True, log_tickers=True,
+                    tag='process_foo__1_print_balances_begin')
+
     def test_parse_args(self):
         processor = market.Processor(self.m)