4 files changed, 108 insertions, 20 deletions

diff --git a/market.py b/market.py
index 6edd605..5876071 100644
--- a/market.py
+++ b/market.py
@@ -255,6 +255,7 @@ class Processor:
                     "name": "print_balances",
                     "number": 1,
                     "fetch_balances": ["begin"],
+                    "fetch_balances_args": { "add_portfolio": True },
                     "print_tickers": { "base_currency": "BTC" },
                     }
                 ],
@@ -386,15 +387,19 @@ class Processor:
     def process_step(self, scenario_name, step, kwargs):
         process_name = "process_{}__{}_{}".format(scenario_name, step["number"], step["name"])
         self.market.report.log_stage("{}_begin".format(process_name))
+
+        fetch_args = step.get("fetch_balances_args", {})
         if "begin" in step.get("fetch_balances", []):
-            self.market.balances.fetch_balances(tag="{}_begin".format(process_name), log_tickers=True)
+            self.market.balances.fetch_balances(tag="{}_begin".format(process_name),
+                    log_tickers=True, **fetch_args)
 
         for action in self.ordered_actions:
             if action in step:
                 self.run_action(action, step[action], kwargs)
 
         if "end" in step.get("fetch_balances", []):
-            self.market.balances.fetch_balances(tag="{}_end".format(process_name), log_tickers=True)
+            self.market.balances.fetch_balances(tag="{}_end".format(process_name),
+                    log_tickers=True, **fetch_args)
         self.market.report.log_stage("{}_end".format(process_name))
 
     def method_arguments(self, action):
diff --git a/store.py b/store.py
index 81828f0..76cfec8 100644
--- a/store.py
+++ b/store.py
@@ -303,13 +303,16 @@ class BalanceStore:
                 compute_value, type)
         return amounts
 
-    def fetch_balances(self, tag=None, log_tickers=False,
+    def fetch_balances(self, tag=None, add_portfolio=False, log_tickers=False,
             ticker_currency="BTC", ticker_compute_value="average", ticker_type="total"):
         all_balances = self.market.ccxt.fetch_all_balances()
         for currency, balance in all_balances.items():
             if balance["exchange_total"] != 0 or balance["margin_total"] != 0 or \
                     currency in self.all:
                 self.all[currency] = portfolio.Balance(currency, balance)
+        if add_portfolio:
+            for currency in Portfolio.repartition(from_cache=True):
+                self.all.setdefault(currency, portfolio.Balance(currency, {}))
         if log_tickers:
             tickers = self.in_currency(ticker_currency, compute_value=ticker_compute_value, type=ticker_type)
             self.market.report.log_balances(tag=tag,
@@ -509,7 +512,9 @@ class Portfolio:
             cls.get_cryptoportfolio(refetch=True)
 
     @classmethod
-    def repartition(cls, liquidity="medium"):
+    def repartition(cls, liquidity="medium", from_cache=False):
+        if from_cache:
+            cls.retrieve_cryptoportfolio()
         cls.get_cryptoportfolio()
         liquidities = cls.liquidities.get(liquidity)
         return liquidities[cls.last_date.get()]
@@ -538,6 +543,20 @@ class Portfolio:
             cls.liquidities.set({})
 
     @classmethod
+    def retrieve_cryptoportfolio(cls):
+        if dbs.redis_connected():
+            repartition = dbs.redis.get("/cryptoportfolio/repartition/latest")
+            date = dbs.redis.get("/cryptoportfolio/repartition/date")
+            if date is not None and repartition is not None:
+                date = datetime.datetime.strptime(date.decode(), "%Y-%m-%d")
+                repartition = json.loads(repartition, parse_int=D, parse_float=D)
+                repartition = { k: { date: v } for k, v in repartition.items() }
+
+                cls.data.set("")
+                cls.last_date.set(date)
+                cls.liquidities.set(repartition)
+
+    @classmethod
     def store_cryptoportfolio(cls):
         if dbs.redis_connected():
             hash_ = {}
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")