Update python tests

nautechsystems · Oct 23, 2023 · 477d0fa · 477d0fa
1 parent dcea95d
commit 477d0fa
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 18 deletions.
diff --git a/nautilus_core/network/src/lib.rs b/nautilus_core/network/src/lib.rs
@@ -23,7 +23,7 @@ use http::{HttpClient, HttpMethod, HttpResponse};
 use pyo3::prelude::*;
 use ratelimiter::quota::Quota;
 use socket::{SocketClient, SocketConfig};
-use websocket::WebSocketClient;
+use websocket::{WebSocketClient, WebSocketConfig};
 
 /// Loaded as nautilus_pyo3.network
 #[pymodule]
@@ -33,6 +33,7 @@ pub fn network(_: Python<'_>, m: &PyModule) -> PyResult<()> {
     m.add_class::<Quota>()?;
     m.add_class::<HttpResponse>()?;
     m.add_class::<WebSocketClient>()?;
+    m.add_class::<WebSocketConfig>()?;
     m.add_class::<SocketClient>()?;
     m.add_class::<SocketConfig>()?;
     Ok(())

diff --git a/tests/integration_tests/network/test_websocket.py b/tests/integration_tests/network/test_websocket.py
@@ -20,6 +20,7 @@
 from aiohttp.test_utils import TestServer
 
 from nautilus_trader.core.nautilus_pyo3 import WebSocketClient
+from nautilus_trader.core.nautilus_pyo3 import WebSocketConfig
 from nautilus_trader.test_kit.functions import eventually
 
 
@@ -31,11 +32,8 @@ def _server_url(server: TestServer) -> str:
 async def test_connect_and_disconnect(websocket_server):
     # Arrange
     store = []
-
-    client = await WebSocketClient.connect(
-        url=_server_url(websocket_server),
-        handler=store.append,
-    )
+    config = WebSocketConfig(_server_url(websocket_server), store.append, [])
+    client = await WebSocketClient.connect(config)
 
     # Act, Assert
     await eventually(lambda: client.is_alive)
@@ -47,10 +45,8 @@ async def test_connect_and_disconnect(websocket_server):
 async def test_client_send_recv(websocket_server):
     # Arrange
     store = []
-    client = await WebSocketClient.connect(
-        url=_server_url(websocket_server),
-        handler=store.append,
-    )
+    config = WebSocketConfig(_server_url(websocket_server), store.append, [])
+    client = await WebSocketClient.connect(config)
     await eventually(lambda: client.is_alive)
 
     # Act
@@ -69,10 +65,8 @@ async def test_client_send_recv(websocket_server):
 async def test_client_send_recv_json(websocket_server):
     # Arrange
     store = []
-    client = await WebSocketClient.connect(
-        url=_server_url(websocket_server),
-        handler=store.append,
-    )
+    config = WebSocketConfig(_server_url(websocket_server), store.append, [])
+    client = await WebSocketClient.connect(config)
     await eventually(lambda: client.is_alive)
 
     # Act
@@ -92,10 +86,8 @@ async def test_client_send_recv_json(websocket_server):
 async def test_reconnect_after_close(websocket_server):
     # Arrange
     store = []
-    client = await WebSocketClient.connect(
-        url=_server_url(websocket_server),
-        handler=store.append,
-    )
+    config = WebSocketConfig(_server_url(websocket_server), store.append, [])
+    client = await WebSocketClient.connect(config)
     await eventually(lambda: client.is_alive)
 
     # Act