1 files changed, 142 insertions, 0 deletions

diff --git a/test.py b/test.py
index 21077a3..d10bfe4 100644
--- a/test.py
+++ b/test.py
@@ -251,5 +251,147 @@ class PortfolioTest(unittest.TestCase):
     def tearDown(self):
         self.patcher.stop()
 
+class BalanceTest(unittest.TestCase):
+    def setUp(self):
+        super(BalanceTest, self).setUp()
+
+        self.fetch_balance = {
+                "free": "foo",
+                "info": "bar",
+                "used": "baz",
+                "total": "bazz",
+                "USDT": {
+                    "free": 6.0,
+                    "used": 1.2,
+                    "total": 7.2
+                    },
+                "XVG": {
+                    "free": 16,
+                    "used": 0.0,
+                    "total": 16
+                    },
+                "XMR": {
+                    "free": 0.0,
+                    "used": 0.0,
+                    "total": 0.0
+                    },
+                }
+        self.patcher = mock.patch.multiple(portfolio.Balance, known_balances={}, trades={})
+        self.patcher.start()
+
+    def test_values(self):
+        balance = portfolio.Balance("BTC", 0.65, 0.35, 0.30)
+        self.assertEqual(0.65, balance.total.value)
+        self.assertEqual(0.35, balance.free.value)
+        self.assertEqual(0.30, balance.used.value)
+        self.assertEqual("BTC", balance.currency)
+
+        balance = portfolio.Balance.from_hash("BTC", { "total": 0.65, "free": 0.35, "used": 0.30})
+        self.assertEqual(0.65, balance.total.value)
+        self.assertEqual(0.35, balance.free.value)
+        self.assertEqual(0.30, balance.used.value)
+        self.assertEqual("BTC", balance.currency)
+
+    @mock.patch.object(portfolio.Amount, "get_ticker")
+    def test_in_currency(self, get_ticker):
+        portfolio.Balance.known_balances = {
+                "BTC": portfolio.Balance("BTC", 0.65, 0.35, 0.30),
+                "ETH": portfolio.Balance("ETH", 3, 3, 0),
+                }
+        market = mock.Mock()
+        get_ticker.return_value = {
+                "bid": 0.09,
+                "ask": 0.11,
+                "average": 0.1,
+                }
+
+        amounts = portfolio.Balance.in_currency("BTC", market)
+        self.assertEqual("BTC", amounts["ETH"].currency)
+        self.assertEqual(0.65, amounts["BTC"].value)
+        self.assertEqual(0.30, amounts["ETH"].value)
+
+        amounts = portfolio.Balance.in_currency("BTC", market, action="bid")
+        self.assertEqual(0.65, amounts["BTC"].value)
+        self.assertEqual(0.27, amounts["ETH"].value)
+
+        amounts = portfolio.Balance.in_currency("BTC", market, action="bid", type="used")
+        self.assertEqual(0.30, amounts["BTC"].value)
+        self.assertEqual(0, amounts["ETH"].value)
+
+    def test_currencies(self):
+        portfolio.Balance.known_balances = {
+                "BTC": portfolio.Balance("BTC", 0.65, 0.35, 0.30),
+                "ETH": portfolio.Balance("ETH", 3, 3, 0),
+                }
+        self.assertListEqual(["BTC", "ETH"], list(portfolio.Balance.currencies()))
+
+    @mock.patch.object(portfolio.market, "fetch_balance")
+    def test_fetch_balances(self, fetch_balance):
+        fetch_balance.return_value = self.fetch_balance
+
+        portfolio.Balance.fetch_balances(portfolio.market)
+        self.assertNotIn("XMR", portfolio.Balance.currencies())
+        self.assertEqual(["USDT", "XVG"], list(portfolio.Balance.currencies()))
+
+    @mock.patch.object(portfolio.Portfolio, "repartition_pertenthousand")
+    @mock.patch.object(portfolio.market, "fetch_balance")
+    def test_dispatch_assets(self, fetch_balance, repartition):
+        fetch_balance.return_value = self.fetch_balance
+        portfolio.Balance.fetch_balances(portfolio.market)
+
+        self.assertNotIn("XEM", portfolio.Balance.currencies())
+
+        repartition.return_value = {
+                "XEM": 7500,
+                "BTC": 2600,
+                }
+
+        amounts = portfolio.Balance.dispatch_assets(portfolio.Amount("BTC", 10.1))
+        self.assertIn("XEM", portfolio.Balance.currencies())
+        self.assertEqual(2.6, amounts["BTC"].value)
+        self.assertEqual(7.5, amounts["XEM"].value)
+
+    @mock.patch.object(portfolio.Portfolio, "repartition_pertenthousand")
+    @mock.patch.object(portfolio.Amount, "get_ticker")
+    @mock.patch.object(portfolio.Trade, "compute_trades")
+    def test_prepare_trades(self, compute_trades, get_ticker, repartition):
+        repartition.return_value = {
+                "XEM": 7500,
+                "BTC": 2500,
+                }
+        get_ticker.side_effect = [
+                { "average": 0.0001 },
+                { "average": 0.000001 }
+                ]
+        market = mock.Mock()
+        market.fetch_balance.return_value = {
+                "USDT": {
+                    "free": 10000.0,
+                    "used": 0.0,
+                    "total": 10000.0
+                    },
+                "XVG": {
+                    "free": 10000.0,
+                    "used": 0.0,
+                    "total": 10000.0
+                    },
+                }
+        portfolio.Balance.prepare_trades(market)
+        compute_trades.assert_called()
+
+        call = compute_trades.call_args
+        self.assertEqual(market, call[1]["market"])
+        self.assertEqual(1, call[0][0]["USDT"].value)
+        self.assertEqual(0.01, call[0][0]["XVG"].value)
+        self.assertEqual(0.2525, call[0][1]["BTC"].value)
+        self.assertEqual(0.7575, call[0][1]["XEM"].value)
+
+    def test__repr(self):
+        balance = portfolio.Balance("BTX", 3, 1, 2)
+        self.assertEqual("Balance(BTX [1.00000000 BTX/2.00000000 BTX/3.00000000 BTX])", repr(balance))
+
+    def tearDown(self):
+        self.patcher.stop()
+
 if __name__ == '__main__':
     unittest.main()