Include current portfolio currencies when printing balances

2 files changed, 80 insertions, 16 deletions

diff --git a/tests/test_market.py b/tests/test_market.py
index ab3cd5e..46fad53 100644
--- a/tests/test_market.py
+++ b/tests/test_market.py
@@ -1032,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)
 
diff --git a/tests/test_store.py b/tests/test_store.py
index d0f7755..ee7e063 100644
--- a/tests/test_store.py
+++ b/tests/test_store.py
@@ -391,6 +391,18 @@ class BalanceStoreTest(WebMockTestCase):
                     tag=None, ticker_currency='FOO', tickers='tickers',
                     type='type')
 
+        balance_store = market.BalanceStore(self.m)
+
+        with self.subTest(add_portfolio=True),\
+                mock.patch.object(market.Portfolio, "repartition") as repartition:
+            repartition.return_value = {
+                    "DOGE": D("0.5"),
+                    "USDT": D("0.5"),
+                    }
+            balance_store.fetch_balances(add_portfolio=True)
+            self.assertListEqual(["USDT", "XVG", "XMR", "DOGE"], list(balance_store.currencies()))
+
+
     @mock.patch.object(market.Portfolio, "repartition")
     def test_dispatch_assets(self, repartition):
         self.m.ccxt.fetch_all_balances.return_value = self.fetch_balance
@@ -1259,24 +1271,67 @@ class PortfolioTest(WebMockTestCase):
                 mock.call("/cryptoportfolio/repartition/date", "2018-03-08"),
                 ])
 
-    @mock.patch.object(market.Portfolio, "get_cryptoportfolio")
-    def test_repartition(self, get_cryptoportfolio):
-        market.Portfolio.liquidities = store.LockedVar({
-                "medium": {
-                    "2018-03-01": "medium_2018-03-01",
-                    "2018-03-08": "medium_2018-03-08",
-                    },
-                "high": {
-                    "2018-03-01": "high_2018-03-01",
-                    "2018-03-08": "high_2018-03-08",
+    @mock.patch.object(store.dbs, "redis_connected")
+    @mock.patch.object(store.dbs, "redis")
+    def test_retrieve_cryptoportfolio(self, redis, redis_connected):
+        with self.subTest(redis_connected=False):
+            redis_connected.return_value = False
+            store.Portfolio.retrieve_cryptoportfolio()
+            redis.get.assert_not_called()
+            self.assertIsNone(store.Portfolio.data.get())
+
+        with self.subTest(redis_connected=True, value=None):
+            redis_connected.return_value = True
+            redis.get.return_value = None
+            store.Portfolio.retrieve_cryptoportfolio()
+            self.assertEqual(2, redis.get.call_count)
+
+        redis.reset_mock()
+        with self.subTest(redis_connected=True, value="present"):
+            redis_connected.return_value = True
+            redis.get.side_effect = [
+                    b'{ "medium": "medium_repartition", "high": "high_repartition" }',
+                    b"2018-03-08"
+                    ]
+            store.Portfolio.retrieve_cryptoportfolio()
+            self.assertEqual(2, redis.get.call_count)
+            self.assertEqual(datetime.datetime(2018,3,8), store.Portfolio.last_date.get())
+            self.assertEqual("", store.Portfolio.data.get())
+            expected_liquidities = {
+                    'high': { datetime.datetime(2018, 3, 8): 'high_repartition' },
+                    'medium': { datetime.datetime(2018, 3, 8): 'medium_repartition' },
                     }
-                })
-        market.Portfolio.last_date = store.LockedVar("2018-03-08")
+            self.assertEqual(expected_liquidities, store.Portfolio.liquidities.get())
+
+    @mock.patch.object(market.Portfolio, "get_cryptoportfolio")
+    @mock.patch.object(market.Portfolio, "retrieve_cryptoportfolio")
+    def test_repartition(self, retrieve_cryptoportfolio, get_cryptoportfolio):
+        with self.subTest(from_cache=False):
+            market.Portfolio.liquidities = store.LockedVar({
+                    "medium": {
+                        "2018-03-01": "medium_2018-03-01",
+                        "2018-03-08": "medium_2018-03-08",
+                        },
+                    "high": {
+                        "2018-03-01": "high_2018-03-01",
+                        "2018-03-08": "high_2018-03-08",
+                        }
+                    })
+            market.Portfolio.last_date = store.LockedVar("2018-03-08")
+
+            self.assertEqual("medium_2018-03-08", market.Portfolio.repartition())
+            get_cryptoportfolio.assert_called_once_with()
+            retrieve_cryptoportfolio.assert_not_called()
+            self.assertEqual("medium_2018-03-08", market.Portfolio.repartition(liquidity="medium"))
+            self.assertEqual("high_2018-03-08", market.Portfolio.repartition(liquidity="high"))
+
+        retrieve_cryptoportfolio.reset_mock()
+        get_cryptoportfolio.reset_mock()
 
-        self.assertEqual("medium_2018-03-08", market.Portfolio.repartition())
-        get_cryptoportfolio.assert_called_once_with()
-        self.assertEqual("medium_2018-03-08", market.Portfolio.repartition(liquidity="medium"))
-        self.assertEqual("high_2018-03-08", market.Portfolio.repartition(liquidity="high"))
+        with self.subTest(from_cache=True):
+            self.assertEqual("medium_2018-03-08", market.Portfolio.repartition(from_cache=True))
+            get_cryptoportfolio.assert_called_once_with()
+            retrieve_cryptoportfolio.assert_called_once_with()
 
     @mock.patch.object(market.time, "sleep")
     @mock.patch.object(market.Portfolio, "get_cryptoportfolio")