diff --git a/bfxapi/tests/helpers.py b/bfxapi/tests/helpers.py index e1b5fee..6edcf66 100644 --- a/bfxapi/tests/helpers.py +++ b/bfxapi/tests/helpers.py @@ -2,23 +2,33 @@ import time import json import asyncio -from .. import Client, BfxWebsocket +from .. import Client, BfxWebsocket, Socket def get_now(): return int(round(time.time() * 1000)) +def ev_worker_override(): + return asyncio.get_event_loop() + class StubbedWebsocket(BfxWebsocket): - def __new__(cls, *args, **kwargs): - instance = super(StubbedWebsocket, cls).__new__(cls, *args, **kwargs) - instance.sent_items = [] - instance.published_items = [] - return instance + + def __init__(self, *args, **kwargs): + self.sent_items = [] + self.published_items = [] + super().__init__(create_event_emitter=ev_worker_override, *args, **kwargs) async def _main(self, host): print ("Faking wesocket connection to {}".format(host)) - def get_ws(self): - return self + def _start_new_socket(self): + socket = Socket(len(self.sockets)) + socket.set_connected() + socket.ws = self + self.sockets[socket.id] = socket + return socket.id + + def _wait_for_socket(self, socketId): + return async def publish(self, data, is_json=True): self.published_items += [{ @@ -27,7 +37,7 @@ class StubbedWebsocket(BfxWebsocket): }] # convert to string and push through the websocket data = json.dumps(data) if is_json else data - return await self.on_message(data) + return await self.on_message(0, data) async def publish_auth_confirmation(self): return self.publish({"event":"auth","status":"OK","chanId":0,"userId":269499,"auth_id":"58aa0472-b1a9-4690-8ab8-300d68e66aaf","caps":{"orders":{"read":1,"write":1},"account":{"read":1,"write":0},"funding":{"read":1,"write":1},"history":{"read":1,"write":0},"wallets":{"read":1,"write":1},"withdraw":{"read":0,"write":1},"positions":{"read":1,"write":1}}}) @@ -73,11 +83,24 @@ class EventWatcher(): counter += 1 return self.value +class StubClient(): + ws = None + res = None + def create_stubbed_client(*args, **kwargs): - client = Client(*args, **kwargs) + client = StubClient() # no support for rest stubbing yet client.rest = None - client.ws = StubbedWebsocket(*args, **kwargs) + wsStub = StubbedWebsocket(*args, **kwargs) + # stub client.ws so tests can use publish + client.ws = wsStub + client.ws.API_KEY = "test key" + client.ws.API_SECRET = "secret key" + # stub socket so we can track socket send requests + socket = Socket(0) + socket.set_connected() + socket.ws = wsStub + client.ws.sockets = { 0: socket } return client async def ws_publish_auth_accepted(ws): diff --git a/bfxapi/tests/test_ws_capacity.py b/bfxapi/tests/test_ws_capacity.py new file mode 100644 index 0000000..d7b28fb --- /dev/null +++ b/bfxapi/tests/test_ws_capacity.py @@ -0,0 +1,42 @@ +import pytest +import json +import time +import asyncio +from .helpers import (create_stubbed_client, ws_publish_connection_init, ws_publish_auth_accepted) + +@pytest.mark.asyncio +async def test_ws_creates_new_socket(): + client = create_stubbed_client() + client.ws.ws_capacity = 5 + # publsh connection created message + await ws_publish_connection_init(client.ws) + # create a bunch of websocket subscriptions + for symbol in ['tXRPBTC', 'tLTCUSD']: + await client.ws.subscribe('candles', symbol, timeframe='1m') + assert len(client.ws.sockets) == 1 + assert client.ws.get_total_available_capcity() == 3 + # subscribe to a few more to force the lib to create a new ws conenction + for symbol in ['tETHBTC', 'tBTCUSD', 'tETHUSD', 'tLTCBTC']: + await client.ws.subscribe('candles', symbol, timeframe='1m') + assert len(client.ws.sockets) == 2 + assert client.ws.get_total_available_capcity() == 4 + +@pytest.mark.asyncio +async def test_ws_uses_authenticated_socket(): + client = create_stubbed_client() + client.ws.ws_capacity = 2 + # publsh connection created message + await ws_publish_connection_init(client.ws) + # create a bunch of websocket subscriptions + for symbol in ['tXRPBTC', 'tLTCUSD', 'tETHBTC', 'tBTCUSD', 'tETHUSD', 'tLTCBTC']: + await client.ws.subscribe('candles', symbol, timeframe='1m') + # publsh connection created message on socket (0 by default) + await ws_publish_connection_init(client.ws) + # send auth accepted (on socket by default) + await ws_publish_auth_accepted(client.ws) + # socket 0 should be the authenticated socket + assert client.ws.get_authenticated_socket().id == 0 + # there should be no other authenticated sockets + for socket in client.ws.sockets.values(): + if socket.id != 0: + assert socket.isAuthenticated == False diff --git a/bfxapi/tests/test_ws_orders.py b/bfxapi/tests/test_ws_orders.py index 39067e8..ebf5e5b 100644 --- a/bfxapi/tests/test_ws_orders.py +++ b/bfxapi/tests/test_ws_orders.py @@ -1,5 +1,6 @@ import pytest import json +import asyncio from .helpers import (create_stubbed_client, ws_publish_auth_accepted, ws_publish_connection_init, EventWatcher) diff --git a/bfxapi/utils/testing_tools.py b/bfxapi/utils/testing_tools.py index 143df5d..121ff51 100644 --- a/bfxapi/utils/testing_tools.py +++ b/bfxapi/utils/testing_tools.py @@ -17,9 +17,6 @@ class StubbedWebsocket(BfxWebsocket): async def _main(self, host): print ("Faking wesocket connection to {}".format(host)) - def get_ws(self): - return self - async def publish(self, data, is_json=True): self.published_items += [{ 'time': get_now(),