tests: add websocket capacity test

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):

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

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)
 

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(),