nautilus_network/python/
websocket.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16use std::{
17    sync::{
18        Arc,
19        atomic::{AtomicU8, Ordering},
20    },
21    time::Duration,
22};
23
24use nautilus_core::{
25    collections::into_ustr_vec,
26    python::{clone_py_object, to_pyruntime_err, to_pyvalue_err},
27};
28use pyo3::{Py, create_exception, exceptions::PyException, prelude::*, types::PyBytes};
29
30use crate::{
31    RECONNECTED,
32    mode::ConnectionMode,
33    ratelimiter::quota::Quota,
34    transport::{Message, TransportError},
35    websocket::{
36        TransportBackend, WebSocketClient, WebSocketConfig,
37        types::{MessageHandler, PingHandler, WriterCommand},
38    },
39};
40
41create_exception!(network, WebSocketClientError, PyException);
42
43#[expect(clippy::needless_pass_by_value)]
44fn to_websocket_pyerr(e: TransportError) -> PyErr {
45    PyErr::new::<WebSocketClientError, _>(e.to_string())
46}
47
48fn is_python_reconnect_control_message(msg: &Message) -> bool {
49    matches!(msg, Message::Text(text) if text.as_ref() == RECONNECTED.as_bytes())
50}
51
52#[pymethods]
53#[pyo3_stub_gen::derive::gen_stub_pymethods]
54impl WebSocketConfig {
55    /// Configuration for WebSocket client connections.
56    ///
57    /// This struct contains only static configuration settings. Runtime callbacks
58    /// (message handler, ping handler) are passed separately to `connect()`.
59    ///
60    /// # Connection Modes
61    ///
62    /// ## Handler Mode
63    ///
64    /// - Use with `crate.websocket.WebSocketClient.connect`.
65    /// - Pass a message handler to `connect()` to receive messages via callback.
66    /// - Client spawns internal task to read messages and call handler.
67    /// - Supports automatic reconnection with exponential backoff.
68    /// - Reconnection config fields (`reconnect_*`) are active.
69    /// - Best for long-lived connections, Python bindings, callback-based APIs.
70    ///
71    /// ## Stream Mode
72    ///
73    /// - Use with `crate.websocket.WebSocketClient.connect_stream`.
74    /// - Returns a `MessageReader` stream for the caller to read from.
75    /// - **Does NOT support automatic reconnection** (reader owned by caller).
76    /// - Reconnection config fields are ignored.
77    /// - On disconnect, client transitions to CLOSED state and caller must manually reconnect.
78    #[new]
79    #[expect(clippy::too_many_arguments)]
80    #[pyo3(signature = (
81        url,
82        headers,
83        heartbeat=None,
84        heartbeat_msg=None,
85        reconnect_timeout_ms=10_000,
86        reconnect_delay_initial_ms=2_000,
87        reconnect_delay_max_ms=30_000,
88        reconnect_backoff_factor=1.5,
89        reconnect_jitter_ms=100,
90        reconnect_max_attempts=None,
91        idle_timeout_ms=None,
92        proxy_url=None,
93    ))]
94    fn py_new(
95        url: String,
96        headers: Vec<(String, String)>,
97        heartbeat: Option<u64>,
98        heartbeat_msg: Option<String>,
99        reconnect_timeout_ms: Option<u64>,
100        reconnect_delay_initial_ms: Option<u64>,
101        reconnect_delay_max_ms: Option<u64>,
102        reconnect_backoff_factor: Option<f64>,
103        reconnect_jitter_ms: Option<u64>,
104        reconnect_max_attempts: Option<u32>,
105        idle_timeout_ms: Option<u64>,
106        proxy_url: Option<String>,
107    ) -> PyResult<Self> {
108        let config = Self {
109            url,
110            headers,
111            heartbeat,
112            heartbeat_msg,
113            reconnect_timeout_ms,
114            reconnect_delay_initial_ms,
115            reconnect_delay_max_ms,
116            reconnect_backoff_factor,
117            reconnect_jitter_ms,
118            reconnect_max_attempts,
119            idle_timeout_ms,
120            backend: TransportBackend::default(),
121            proxy_url,
122        };
123        config.validate().map_err(to_pyvalue_err)?;
124        Ok(config)
125    }
126}
127
128#[pymethods]
129#[pyo3_stub_gen::derive::gen_stub_pymethods]
130impl WebSocketClient {
131    /// Creates a websocket client in **handler mode** with automatic reconnection.
132    ///
133    /// The handler is called for each incoming message on an internal task.
134    /// Automatic reconnection is **enabled** with exponential backoff. On disconnection,
135    /// the client automatically attempts to reconnect and replaces the internal reader
136    /// (the handler continues working seamlessly).
137    ///
138    /// Use handler mode for simplified connection management, automatic reconnection, Python
139    /// bindings, or callback-based message handling.
140    ///
141    /// See `WebSocketConfig` documentation for comparison with stream mode.
142    #[staticmethod]
143    #[pyo3(name = "connect", signature = (loop_, config, handler, ping_handler = None, post_reconnection = None, keyed_quotas = Vec::new(), default_quota = None))]
144    #[expect(clippy::too_many_arguments, clippy::needless_pass_by_value)]
145    fn py_connect(
146        loop_: Py<PyAny>,
147        config: WebSocketConfig,
148        handler: Py<PyAny>,
149        ping_handler: Option<Py<PyAny>>,
150        post_reconnection: Option<Py<PyAny>>,
151        keyed_quotas: Vec<(String, Quota)>,
152        default_quota: Option<Quota>,
153        py: Python<'_>,
154    ) -> PyResult<Bound<'_, PyAny>> {
155        let call_soon_threadsafe: Py<PyAny> = loop_.getattr(py, "call_soon_threadsafe")?;
156        let call_soon_clone = clone_py_object(&call_soon_threadsafe);
157        let handler_clone = clone_py_object(&handler);
158
159        let message_handler: MessageHandler = Arc::new(move |msg: Message| {
160            if is_python_reconnect_control_message(&msg) {
161                return;
162            }
163
164            Python::attach(|py| {
165                let py_bytes = match &msg {
166                    Message::Binary(data) | Message::Text(data) => PyBytes::new(py, data.as_ref()),
167                    _ => return,
168                };
169
170                if let Err(e) = call_soon_clone.call1(py, (&handler_clone, py_bytes)) {
171                    log::error!("Error scheduling message handler on event loop: {e}");
172                }
173            });
174        });
175
176        let ping_handler_fn = ping_handler.map(|ping_handler| {
177            let ping_handler_clone = clone_py_object(&ping_handler);
178            let call_soon_clone = clone_py_object(&call_soon_threadsafe);
179
180            let ping_handler_fn: PingHandler = Arc::new(move |data: Vec<u8>| {
181                Python::attach(|py| {
182                    let py_bytes = PyBytes::new(py, &data);
183                    if let Err(e) = call_soon_clone.call1(py, (&ping_handler_clone, py_bytes)) {
184                        log::error!("Error scheduling ping handler on event loop: {e}");
185                    }
186                });
187            });
188            ping_handler_fn
189        });
190
191        let post_reconnection_fn = post_reconnection.map(|callback| {
192            let callback_clone = clone_py_object(&callback);
193            Arc::new(move || {
194                Python::attach(|py| {
195                    if let Err(e) = callback_clone.call0(py) {
196                        log::error!("Error calling post_reconnection handler: {e}");
197                    }
198                });
199            }) as std::sync::Arc<dyn Fn() + Send + Sync>
200        });
201
202        pyo3_async_runtimes::tokio::future_into_py(py, async move {
203            Box::pin(Self::connect(
204                config,
205                Some(message_handler),
206                ping_handler_fn,
207                post_reconnection_fn,
208                keyed_quotas,
209                default_quota,
210            ))
211            .await
212            .map_err(to_websocket_pyerr)
213        })
214    }
215
216    /// Set disconnect mode to true.
217    ///
218    /// Controller task will periodically check the disconnect mode
219    /// and shutdown the client if it is alive
220    ///
221    /// If an `AuthTracker` is registered, this fails pending auth waits.
222    #[pyo3(name = "disconnect")]
223    #[expect(clippy::needless_pass_by_value)]
224    fn py_disconnect<'py>(slf: PyRef<'_, Self>, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
225        let connection_mode = slf.connection_mode.clone();
226        let state_notify = slf.state_notify.clone();
227        let mode = ConnectionMode::from_atomic(&connection_mode);
228        log::debug!("Close from mode {mode}");
229
230        pyo3_async_runtimes::tokio::future_into_py(py, async move {
231            match ConnectionMode::from_atomic(&connection_mode) {
232                ConnectionMode::Closed => {
233                    log::debug!("WebSocket already closed");
234                }
235                ConnectionMode::Disconnect => {
236                    log::debug!("WebSocket already disconnecting");
237                }
238                _ => {
239                    // Preserve a CLOSED terminal state reached concurrently
240                    ConnectionMode::request_disconnect(&connection_mode);
241                    state_notify.notify_one();
242
243                    let timeout = tokio::time::timeout(Duration::from_secs(5), async {
244                        while !ConnectionMode::from_atomic(&connection_mode).is_closed() {
245                            tokio::time::sleep(Duration::from_millis(10)).await;
246                        }
247                    })
248                    .await;
249
250                    if timeout.is_err() {
251                        log::warn!("Timeout waiting for WebSocket to close, forcing closed state");
252                        connection_mode.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
253                    }
254                }
255            }
256
257            Ok(())
258        })
259    }
260
261    /// Check if the client connection is active.
262    ///
263    /// Returns `true` if the client is connected and has not been signalled to disconnect.
264    /// The client will automatically retry connection based on its configuration.
265    #[pyo3(name = "is_active")]
266    #[expect(clippy::needless_pass_by_value)]
267    fn py_is_active(slf: PyRef<'_, Self>) -> bool {
268        !slf.controller_task.is_finished()
269    }
270
271    /// Check if the client is reconnecting.
272    ///
273    /// Returns `true` if the client lost connection and is attempting to reestablish it.
274    /// The client will automatically retry connection based on its configuration.
275    #[pyo3(name = "is_reconnecting")]
276    #[expect(clippy::needless_pass_by_value)]
277    fn py_is_reconnecting(slf: PyRef<'_, Self>) -> bool {
278        slf.is_reconnecting()
279    }
280
281    /// Check if the client is disconnecting.
282    ///
283    /// Returns `true` if the client is in disconnect mode.
284    #[pyo3(name = "is_disconnecting")]
285    #[expect(clippy::needless_pass_by_value)]
286    fn py_is_disconnecting(slf: PyRef<'_, Self>) -> bool {
287        slf.is_disconnecting()
288    }
289
290    /// Check if the client is closed.
291    ///
292    /// Returns `true` if the client has been explicitly disconnected or reached
293    /// maximum reconnection attempts. In this state, the client cannot be reused
294    /// and a new client must be created for further connections.
295    #[pyo3(name = "is_closed")]
296    #[expect(clippy::needless_pass_by_value)]
297    fn py_is_closed(slf: PyRef<'_, Self>) -> bool {
298        slf.is_closed()
299    }
300
301    /// Send bytes data to the server.
302    ///
303    /// # Errors
304    ///
305    /// Returns an error if:
306    /// - The connection is not active or closes while waiting for rate limit (`WebSocketClientError`).
307    /// - The writer channel is broken (`PyRuntimeError`).
308    #[pyo3(name = "send")]
309    #[pyo3(signature = (data, keys=None))]
310    #[expect(clippy::needless_pass_by_value)]
311    fn py_send<'py>(
312        slf: PyRef<'_, Self>,
313        data: Vec<u8>,
314        py: Python<'py>,
315        keys: Option<Vec<String>>,
316    ) -> PyResult<Bound<'py, PyAny>> {
317        let rate_limiter = slf.rate_limiter.clone();
318        let writer_tx = slf.writer_tx.clone();
319        let mode = slf.connection_mode.clone();
320        let keys = keys.map(into_ustr_vec);
321
322        pyo3_async_runtimes::tokio::future_into_py(py, async move {
323            if !ConnectionMode::from_atomic(&mode).is_active() {
324                let msg = "Cannot send data: connection not active".to_string();
325                log::warn!("{msg}");
326                return Err(to_websocket_pyerr(TransportError::Io(std::io::Error::new(
327                    std::io::ErrorKind::NotConnected,
328                    msg,
329                ))));
330            }
331
332            tokio::select! {
333                biased;
334                () = rate_limiter.await_keys_ready(keys.as_deref()) => {}
335                () = poll_until_closed(&mode) => {
336                    return Err(to_websocket_pyerr(TransportError::Io(std::io::Error::new(
337                        std::io::ErrorKind::ConnectionAborted,
338                        "Connection closed while waiting for rate limit",
339                    ))));
340                }
341            }
342
343            log::trace!("Sending binary: {data:?}");
344
345            let msg = Message::Binary(data.into());
346            writer_tx
347                .send(WriterCommand::Send(msg))
348                .map_err(to_pyruntime_err)
349        })
350    }
351
352    /// Sends the given text `data` to the server.
353    ///
354    /// Returns `Ok(())` when the message is enqueued to the writer channel. This does NOT
355    /// guarantee delivery: if a disconnect occurs concurrently, the writer task may drop the
356    /// message. During reconnection, messages are buffered and replayed on the new connection.
357    #[pyo3(name = "send_text")]
358    #[pyo3(signature = (data, keys=None))]
359    #[expect(clippy::needless_pass_by_value)]
360    fn py_send_text<'py>(
361        slf: PyRef<'_, Self>,
362        data: Vec<u8>,
363        py: Python<'py>,
364        keys: Option<Vec<String>>,
365    ) -> PyResult<Bound<'py, PyAny>> {
366        let data_str = String::from_utf8(data).map_err(to_pyvalue_err)?;
367        let rate_limiter = slf.rate_limiter.clone();
368        let writer_tx = slf.writer_tx.clone();
369        let mode = slf.connection_mode.clone();
370        let keys = keys.map(into_ustr_vec);
371
372        pyo3_async_runtimes::tokio::future_into_py(py, async move {
373            if !ConnectionMode::from_atomic(&mode).is_active() {
374                return Err(to_websocket_pyerr(TransportError::Io(std::io::Error::new(
375                    std::io::ErrorKind::NotConnected,
376                    "Cannot send text: connection not active",
377                ))));
378            }
379
380            tokio::select! {
381                biased;
382                () = rate_limiter.await_keys_ready(keys.as_deref()) => {}
383                () = poll_until_closed(&mode) => {
384                    return Err(to_websocket_pyerr(TransportError::Io(std::io::Error::new(
385                        std::io::ErrorKind::ConnectionAborted,
386                        "Connection closed while waiting for rate limit",
387                    ))));
388                }
389            }
390
391            log::trace!("Sending text: {data_str}");
392
393            let msg = Message::Text(data_str.into());
394            writer_tx
395                .send(WriterCommand::Send(msg))
396                .map_err(to_pyruntime_err)
397        })
398    }
399
400    /// Sends a pong frame back to the server.
401    #[pyo3(name = "send_pong")]
402    #[expect(clippy::needless_pass_by_value)]
403    fn py_send_pong<'py>(
404        slf: PyRef<'_, Self>,
405        data: Vec<u8>,
406        py: Python<'py>,
407    ) -> PyResult<Bound<'py, PyAny>> {
408        let writer_tx = slf.writer_tx.clone();
409        let mode = slf.connection_mode.clone();
410        let data_len = data.len();
411
412        pyo3_async_runtimes::tokio::future_into_py(py, async move {
413            if !ConnectionMode::from_atomic(&mode).is_active() {
414                log::debug!("Skipping pong: connection not active");
415                return Ok(());
416            }
417            log::trace!("Sending pong frame ({data_len} bytes)");
418
419            let msg = Message::Pong(data.into());
420            writer_tx
421                .send(WriterCommand::Send(msg))
422                .map_err(to_pyruntime_err)
423        })
424    }
425}
426
427async fn poll_until_closed(mode: &Arc<AtomicU8>) {
428    loop {
429        if matches!(
430            ConnectionMode::from_atomic(mode),
431            ConnectionMode::Disconnect | ConnectionMode::Closed
432        ) {
433            break;
434        }
435
436        tokio::time::sleep(Duration::from_millis(100)).await;
437    }
438}
439
440#[cfg(test)]
441mod control_filter_tests {
442    use bytes::Bytes;
443    use rstest::rstest;
444
445    use super::*;
446
447    #[rstest]
448    #[case::reconnected_control(Message::text(RECONNECTED), true)]
449    #[case::application_text(Message::text("application"), false)]
450    #[case::reconnected_prefix(Message::text(format!("{RECONNECTED}:payload")), false)]
451    #[case::reconnected_binary(Message::Binary(Bytes::from_static(RECONNECTED.as_bytes())), false)]
452    #[case::ping(Message::ping(Bytes::new()), false)]
453    fn python_reconnect_control_filter(#[case] msg: Message, #[case] expected: bool) {
454        assert_eq!(is_python_reconnect_control_message(&msg), expected);
455    }
456}
457
458#[cfg(test)]
459mod py_new_tests {
460    use rstest::rstest;
461
462    use super::*;
463
464    #[rstest]
465    fn test_py_new_rejects_empty_url() {
466        let result = WebSocketConfig::py_new(
467            String::new(),
468            vec![],
469            None,
470            None,
471            None,
472            None,
473            None,
474            None,
475            None,
476            None,
477            None,
478            None,
479        );
480
481        assert!(result.is_err());
482    }
483}
484
485#[cfg(test)]
486#[cfg(not(feature = "turmoil"))]
487#[cfg(target_os = "linux")] // Only run network tests on Linux (CI stability)
488mod tests {
489    use std::ffi::CString;
490
491    use futures_util::{SinkExt, StreamExt};
492    use nautilus_core::python::IntoPyObjectNautilusExt;
493    use pyo3::{prelude::*, types::PyBytes};
494    use tokio::{
495        net::TcpListener,
496        task::{self, JoinHandle},
497        time::{Duration, sleep},
498    };
499    use tokio_tungstenite::{
500        accept_hdr_async,
501        tungstenite::{
502            handshake::server::{self, Callback},
503            http::HeaderValue,
504        },
505    };
506
507    use crate::{
508        transport::Message,
509        websocket::{MessageHandler, WebSocketClient, WebSocketConfig},
510    };
511
512    struct TestServer {
513        task: JoinHandle<()>,
514        port: u16,
515    }
516
517    #[derive(Debug, Clone)]
518    struct TestCallback {
519        key: String,
520        value: HeaderValue,
521    }
522
523    impl Callback for TestCallback {
524        #[expect(clippy::panic_in_result_fn)]
525        fn on_request(
526            self,
527            request: &server::Request,
528            response: server::Response,
529        ) -> Result<server::Response, server::ErrorResponse> {
530            let _ = response;
531            let value = request.headers().get(&self.key);
532            assert!(value.is_some());
533
534            if let Some(value) = request.headers().get(&self.key) {
535                assert_eq!(value, self.value);
536            }
537
538            Ok(response)
539        }
540    }
541
542    impl TestServer {
543        async fn setup(key: String, value: String) -> Self {
544            let server = TcpListener::bind("127.0.0.1:0").await.unwrap();
545            let port = TcpListener::local_addr(&server).unwrap().port();
546
547            let test_call_back = TestCallback {
548                key,
549                value: HeaderValue::from_str(&value).unwrap(),
550            };
551
552            // Set up test server
553            let task = task::spawn(async move {
554                // Keep accepting connections
555                loop {
556                    let (conn, _) = server.accept().await.unwrap();
557                    let mut websocket = accept_hdr_async(conn, test_call_back.clone())
558                        .await
559                        .unwrap();
560
561                    task::spawn(async move {
562                        while let Some(Ok(msg)) = websocket.next().await {
563                            match msg {
564                                tokio_tungstenite::tungstenite::protocol::Message::Text(txt)
565                                    if txt == "close-now" =>
566                                {
567                                    log::debug!("Forcibly closing from server side");
568                                    // This sends a close frame, then stops reading
569                                    let _ = websocket.close(None).await;
570                                    break;
571                                }
572                                // Echo text/binary frames
573                                tokio_tungstenite::tungstenite::protocol::Message::Text(_)
574                                | tokio_tungstenite::tungstenite::protocol::Message::Binary(_) => {
575                                    if websocket.send(msg).await.is_err() {
576                                        break;
577                                    }
578                                }
579                                // If the client closes, we also break
580                                tokio_tungstenite::tungstenite::protocol::Message::Close(
581                                    _frame,
582                                ) => {
583                                    let _ = websocket.close(None).await;
584                                    break;
585                                }
586                                // Ignore pings/pongs
587                                _ => {}
588                            }
589                        }
590                    });
591                }
592            });
593
594            Self { task, port }
595        }
596    }
597
598    impl Drop for TestServer {
599        fn drop(&mut self) {
600            self.task.abort();
601        }
602    }
603
604    fn create_test_handler() -> (Py<PyAny>, Py<PyAny>) {
605        let code_raw = "
606class Counter:
607    def __init__(self):
608        self.count = 0
609        self.check = False
610
611    def handler(self, bytes):
612        msg = bytes.decode()
613        if msg == 'ping':
614            self.count += 1
615        elif msg == 'heartbeat message':
616            self.check = True
617
618    def get_check(self):
619        return self.check
620
621    def get_count(self):
622        return self.count
623
624counter = Counter()
625";
626
627        let code = CString::new(code_raw).unwrap();
628        let filename = CString::new("test".to_string()).unwrap();
629        let module = CString::new("test".to_string()).unwrap();
630        Python::attach(|py| {
631            let pymod = PyModule::from_code(py, &code, &filename, &module).unwrap();
632
633            let counter = pymod.getattr("counter").unwrap().into_py_any_unwrap(py);
634            let handler = counter
635                .getattr(py, "handler")
636                .unwrap()
637                .into_py_any_unwrap(py);
638
639            (counter, handler)
640        })
641    }
642
643    #[tokio::test]
644    async fn basic_client_test() {
645        const N: usize = 10;
646
647        Python::initialize();
648
649        let mut success_count = 0;
650        let header_key = "hello-custom-key".to_string();
651        let header_value = "hello-custom-value".to_string();
652
653        let server = TestServer::setup(header_key.clone(), header_value.clone()).await;
654        let (counter, handler) = create_test_handler();
655
656        let config = WebSocketConfig::py_new(
657            format!("ws://127.0.0.1:{}", server.port),
658            vec![(header_key, header_value)],
659            None,
660            None,
661            None,
662            None,
663            None,
664            None,
665            None,
666            None,
667            None,
668            None,
669        )
670        .unwrap();
671
672        let handler_clone = Python::attach(|py| handler.clone_ref(py));
673
674        let message_handler: MessageHandler = std::sync::Arc::new(move |msg: Message| {
675            Python::attach(|py| {
676                let data = match msg {
677                    Message::Binary(data) | Message::Text(data) => data.to_vec(),
678                    _ => return,
679                };
680                let py_bytes = PyBytes::new(py, &data);
681                if let Err(e) = handler_clone.call1(py, (py_bytes,)) {
682                    log::error!("Error calling handler: {e}");
683                }
684            });
685        });
686
687        let client =
688            WebSocketClient::connect(config, Some(message_handler), None, None, vec![], None)
689                .await
690                .unwrap();
691
692        for _ in 0..N {
693            client.send_bytes(b"ping".to_vec(), None).await.unwrap();
694            success_count += 1;
695        }
696
697        sleep(Duration::from_secs(1)).await;
698        let count_value: usize = Python::attach(|py| {
699            counter
700                .getattr(py, "get_count")
701                .unwrap()
702                .call0(py)
703                .unwrap()
704                .extract(py)
705                .unwrap()
706        });
707        assert_eq!(count_value, success_count);
708
709        // Close the connection => client should reconnect automatically
710        client.send_close_message().await.unwrap();
711
712        // Send messages that increment the count
713        sleep(Duration::from_secs(2)).await;
714
715        for _ in 0..N {
716            client.send_bytes(b"ping".to_vec(), None).await.unwrap();
717            success_count += 1;
718        }
719
720        sleep(Duration::from_secs(1)).await;
721        let count_value: usize = Python::attach(|py| {
722            counter
723                .getattr(py, "get_count")
724                .unwrap()
725                .call0(py)
726                .unwrap()
727                .extract(py)
728                .unwrap()
729        });
730        assert_eq!(count_value, success_count);
731        assert_eq!(success_count, N + N);
732
733        client.disconnect().await;
734        assert!(client.is_disconnected());
735    }
736
737    #[tokio::test]
738    async fn message_ping_test() {
739        Python::initialize();
740
741        let header_key = "hello-custom-key".to_string();
742        let header_value = "hello-custom-value".to_string();
743
744        let (checker, handler) = create_test_handler();
745
746        let server = TestServer::setup(header_key.clone(), header_value.clone()).await;
747        let config = WebSocketConfig::py_new(
748            format!("ws://127.0.0.1:{}", server.port),
749            vec![(header_key, header_value)],
750            Some(1),
751            Some("heartbeat message".to_string()),
752            None,
753            None,
754            None,
755            None,
756            None,
757            None,
758            None,
759            None,
760        )
761        .unwrap();
762
763        let handler_clone = Python::attach(|py| handler.clone_ref(py));
764
765        let message_handler: MessageHandler = std::sync::Arc::new(move |msg: Message| {
766            Python::attach(|py| {
767                let data = match msg {
768                    Message::Binary(data) | Message::Text(data) => data.to_vec(),
769                    _ => return,
770                };
771                let py_bytes = PyBytes::new(py, &data);
772                if let Err(e) = handler_clone.call1(py, (py_bytes,)) {
773                    log::error!("Error calling handler: {e}");
774                }
775            });
776        });
777
778        let client =
779            WebSocketClient::connect(config, Some(message_handler), None, None, vec![], None)
780                .await
781                .unwrap();
782
783        sleep(Duration::from_secs(2)).await;
784        let check_value: bool = Python::attach(|py| {
785            checker
786                .getattr(py, "get_check")
787                .unwrap()
788                .call0(py)
789                .unwrap()
790                .extract(py)
791                .unwrap()
792        });
793        assert!(check_value);
794
795        client.disconnect().await;
796        assert!(client.is_disconnected());
797    }
798}
nautilus_network/python/websocket.rs

nautilus_network/python/
websocket.rs
