2 files changed, 153 insertions, 59 deletions

diff --git a/ccxt_wrapper.py b/ccxt_wrapper.py
index 903bfc4..d37c306 100644
--- a/ccxt_wrapper.py
+++ b/ccxt_wrapper.py
@@ -11,29 +11,6 @@ class poloniexE(poloniex):
     def nanoseconds():
         return int(time.time() * 1000000000)
 
-    def nonce(self):
-        return self.nanoseconds()
-
-    def fetch_balance(self, params={}):
-        self.load_markets()
-        balances = self.privatePostReturnCompleteBalances(self.extend({
-            'account': 'all',
-            }, params))
-        result = {'info': balances}
-        currencies = list(balances.keys())
-        for c in range(0, len(currencies)):
-            id = currencies[c]
-            balance = balances[id]
-            currency = self.common_currency_code(id)
-            account = {
-                    'free': decimal.Decimal(balance['available']),
-                    'used': decimal.Decimal(balance['onOrders']),
-                    'total': decimal.Decimal(0.0),
-                    }
-            account['total'] = self.sum(account['free'], account['used'])
-            result[currency] = account
-        return self.parse_balance(result)
-
     def fetch_margin_balance(self):
         """
         portfolio.market.privatePostGetMarginPosition({"currencyPair": "BTC_DASH"})
@@ -200,31 +177,8 @@ class poloniexE(poloniex):
 
         return all_balances
 
-    def parse_ticker(self, ticker, market=None):
-        timestamp = self.milliseconds()
-        symbol = None
-        if market:
-            symbol = market['symbol']
-        return {
-            'symbol': symbol,
-            'timestamp': timestamp,
-            'datetime': self.iso8601(timestamp),
-            'high': decimal.Decimal(ticker['high24hr']),
-            'low': decimal.Decimal(ticker['low24hr']),
-            'bid': decimal.Decimal(ticker['highestBid']),
-            'ask': decimal.Decimal(ticker['lowestAsk']),
-            'vwap': None,
-            'open': None,
-            'close': None,
-            'first': None,
-            'last': decimal.Decimal(ticker['last']),
-            'change': decimal.Decimal(ticker['percentChange']),
-            'percentage': None,
-            'average': None,
-            'baseVolume': decimal.Decimal(ticker['quoteVolume']),
-            'quoteVolume': decimal.Decimal(ticker['baseVolume']),
-            'info': ticker,
-        }
+    def create_exchange_order(self, symbol, type, side, amount, price=None, params={}):
+        return super(poloniexE, self).create_order(symbol, type, side, amount, price=price, params=params)
 
     def create_margin_order(self, symbol, type, side, amount, price=None, lending_rate=None, params={}):
         if type == 'market':
@@ -254,17 +208,6 @@ class poloniexE(poloniex):
         self.orders[id] = order
         return self.extend({'info': response}, order)
 
-    def create_exchange_order(self, symbol, type, side, amount, price=None, params={}):
-        return super(poloniexE, self).create_order(symbol, type, side, amount, price=price, params=params)
-
-    def create_order(self, symbol, type, side, amount, price=None, account="exchange", lending_rate=None, params={}):
-        if account == "exchange":
-            return self.create_exchange_order(symbol, type, side, amount, price=price, params=params)
-        elif account == "margin":
-            return self.create_margin_order(symbol, type, side, amount, price=price, lending_rate=lending_rate, params=params)
-        else:
-            raise NotImplementedError
-
     def order_precision(self, symbol):
         return 8
 
@@ -324,3 +267,73 @@ class poloniexE(poloniex):
                 "total": decimal.Decimal(summary["totalValue"]),
                 }
 
+    def nonce(self):
+        """
+        Wrapped to allow nonce with other libraries
+        """
+        return self.nanoseconds()
+
+    def fetch_balance(self, params={}):
+        """
+        Wrapped to get decimals
+        """
+        self.load_markets()
+        balances = self.privatePostReturnCompleteBalances(self.extend({
+            'account': 'all',
+            }, params))
+        result = {'info': balances}
+        currencies = list(balances.keys())
+        for c in range(0, len(currencies)):
+            id = currencies[c]
+            balance = balances[id]
+            currency = self.common_currency_code(id)
+            account = {
+                    'free': decimal.Decimal(balance['available']),
+                    'used': decimal.Decimal(balance['onOrders']),
+                    'total': decimal.Decimal(0.0),
+                    }
+            account['total'] = self.sum(account['free'], account['used'])
+            result[currency] = account
+        return self.parse_balance(result)
+
+    def parse_ticker(self, ticker, market=None):
+        """
+        Wrapped to get decimals
+        """
+        timestamp = self.milliseconds()
+        symbol = None
+        if market:
+            symbol = market['symbol']
+        return {
+            'symbol': symbol,
+            'timestamp': timestamp,
+            'datetime': self.iso8601(timestamp),
+            'high': decimal.Decimal(ticker['high24hr']),
+            'low': decimal.Decimal(ticker['low24hr']),
+            'bid': decimal.Decimal(ticker['highestBid']),
+            'ask': decimal.Decimal(ticker['lowestAsk']),
+            'vwap': None,
+            'open': None,
+            'close': None,
+            'first': None,
+            'last': decimal.Decimal(ticker['last']),
+            'change': decimal.Decimal(ticker['percentChange']),
+            'percentage': None,
+            'average': None,
+            'baseVolume': decimal.Decimal(ticker['quoteVolume']),
+            'quoteVolume': decimal.Decimal(ticker['baseVolume']),
+            'info': ticker,
+        }
+
+    def create_order(self, symbol, type, side, amount, price=None, account="exchange", lending_rate=None, params={}):
+        """
+        Wrapped to handle margin and exchange accounts
+        """
+        if account == "exchange":
+            return self.create_exchange_order(symbol, type, side, amount, price=price, params=params)
+        elif account == "margin":
+            return self.create_margin_order(symbol, type, side, amount, price=price, lending_rate=lending_rate, params=params)
+        else:
+            raise NotImplementedError
+
+
diff --git a/test.py b/test.py
index 141c9e0..37b0e88 100644
--- a/test.py
+++ b/test.py
@@ -46,6 +46,87 @@ class WebMockTestCase(unittest.TestCase):
         super(WebMockTestCase, self).tearDown()
 
 @unittest.skipUnless("unit" in limits, "Unit skipped")
+class poloniexETest(unittest.TestCase):
+    def setUp(self):
+        super(poloniexETest, self).setUp()
+        self.wm = requests_mock.Mocker()
+        self.wm.start()
+
+        self.s = market.ccxt.poloniexE()
+
+    def tearDown(self):
+        self.wm.stop()
+        super(poloniexETest, self).tearDown()
+
+    def test_nanoseconds(self):
+        with mock.patch.object(market.ccxt.time, "time") as time:
+            time.return_value = 123456.7890123456
+            self.assertEqual(123456789012345, self.s.nanoseconds())
+
+    def test_nonce(self):
+        with mock.patch.object(market.ccxt.time, "time") as time:
+            time.return_value = 123456.7890123456
+            self.assertEqual(123456789012345, self.s.nonce())
+
+    def test_order_precision(self):
+        self.assertEqual(8, self.s.order_precision("FOO"))
+
+    def test_transfer_balance(self):
+        with self.subTest(success=True),\
+                mock.patch.object(self.s, "privatePostTransferBalance") as t:
+            t.return_value = { "success": 1 }
+            result = self.s.transfer_balance("FOO", 12, "exchange", "margin")
+            t.assert_called_once_with({
+                "currency": "FOO",
+                "amount": 12,
+                "fromAccount": "exchange",
+                "toAccount": "margin",
+                "confirmed": 1
+                })
+            self.assertTrue(result)
+
+        with self.subTest(success=False),\
+                mock.patch.object(self.s, "privatePostTransferBalance") as t:
+            t.return_value = { "success": 0 }
+            self.assertFalse(self.s.transfer_balance("FOO", 12, "exchange", "margin"))
+
+    def test_close_margin_position(self):
+        with mock.patch.object(self.s, "privatePostCloseMarginPosition") as c:
+            self.s.close_margin_position("FOO", "BAR")
+            c.assert_called_with({"currencyPair": "BAR_FOO"})
+
+    def test_tradable_balances(self):
+        with mock.patch.object(self.s, "privatePostReturnTradableBalances") as r:
+            r.return_value = {
+                    "FOO": { "exchange": "12.1234", "margin": "0.0123" },
+                    "BAR": { "exchange": "1", "margin": "0" },
+                    }
+            balances = self.s.tradable_balances()
+            self.assertEqual(["FOO", "BAR"], list(balances.keys()))
+            self.assertEqual(["exchange", "margin"], list(balances["FOO"].keys()))
+            self.assertEqual(D("12.1234"), balances["FOO"]["exchange"])
+            self.assertEqual(["exchange", "margin"], list(balances["BAR"].keys()))
+
+    def test_margin_summary(self):
+        with mock.patch.object(self.s, "privatePostReturnMarginAccountSummary") as r:
+            r.return_value = {
+                    "currentMargin": "1.49680968",
+                    "lendingFees": "0.0000001",
+                    "pl": "0.00008254",
+                    "totalBorrowedValue": "0.00673602",
+                    "totalValue": "0.01000000",
+                    "netValue": "0.01008254",
+                    }
+            expected = {
+                    'current_margin': D('1.49680968'),
+                    'gains': D('0.00008254'),
+                    'lending_fees': D('0.0000001'),
+                    'total': D('0.01000000'),
+                    'total_borrowed': D('0.00673602')
+                    }
+            self.assertEqual(expected, self.s.margin_summary())
+
+@unittest.skipUnless("unit" in limits, "Unit skipped")
 class PortfolioTest(WebMockTestCase):
     def fill_data(self):
         if self.json_response is not None: