diff --git a/tests/write.rs b/tests/write.rs new file mode 100644 index 00000000..aa627cce --- /dev/null +++ b/tests/write.rs @@ -0,0 +1,68 @@ +use std::io::{self, Read, Write}; +use tungstenite::{protocol::WebSocketConfig, Message, WebSocket}; + +/// `Write` impl that records call stats and drops the data. +#[derive(Debug, Default)] +struct MockWrite { + written_bytes: usize, + write_count: usize, + flush_count: usize, +} + +impl Read for MockWrite { + fn read(&mut self, _: &mut [u8]) -> io::Result { + Err(io::Error::new(io::ErrorKind::WouldBlock, "reads not supported")) + } +} +impl Write for MockWrite { + fn write(&mut self, buf: &[u8]) -> io::Result { + self.written_bytes += buf.len(); + self.write_count += 1; + Ok(buf.len()) + } + + fn flush(&mut self) -> io::Result<()> { + self.flush_count += 1; + Ok(()) + } +} + +/// Test for write buffering and flushing behaviour. +#[test] +fn write_flush_behaviour() { + const SEND_ME_LEN: usize = 10; + const BATCH_ME_LEN: usize = 11; + const WRITE_BUFFER_SIZE: usize = 600; + + let mut ws = WebSocket::from_raw_socket( + MockWrite::default(), + tungstenite::protocol::Role::Server, + Some(WebSocketConfig { write_buffer_size: WRITE_BUFFER_SIZE, ..<_>::default() }), + ); + + assert_eq!(ws.get_ref().written_bytes, 0); + assert_eq!(ws.get_ref().write_count, 0); + assert_eq!(ws.get_ref().flush_count, 0); + + // `send` writes & flushes immediately + ws.send(Message::Text("Send me!".into())).unwrap(); + assert_eq!(ws.get_ref().written_bytes, SEND_ME_LEN); + assert_eq!(ws.get_ref().write_count, 1); + assert_eq!(ws.get_ref().flush_count, 1); + + // send a batch of messages + for msg in (0..100).map(|_| Message::Text("Batch me!".into())) { + ws.write(msg).unwrap(); + } + // after 55 writes the out_buffer will exceed write_buffer_size=600 + // and so do a single underlying write (not flushing). + assert_eq!(ws.get_ref().written_bytes, 55 * BATCH_ME_LEN + SEND_ME_LEN); + assert_eq!(ws.get_ref().write_count, 2); + assert_eq!(ws.get_ref().flush_count, 1); + + // flushing will perform a single write for the remaining out_buffer & flush. + ws.flush().unwrap(); + assert_eq!(ws.get_ref().written_bytes, 100 * BATCH_ME_LEN + SEND_ME_LEN); + assert_eq!(ws.get_ref().write_count, 3); + assert_eq!(ws.get_ref().flush_count, 2); +}