diff --git a/nautilus_core/network/src/lib.rs b/nautilus_core/network/src/lib.rs index 9809435ab975..6ee52dd2ec00 100644 --- 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::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; Ok(()) diff --git a/tests/integration_tests/network/test_websocket.py b/tests/integration_tests/network/test_websocket.py index 42e636bb01eb..861cd74d414e 100644 --- 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