From a711c51ea2eb141361a6dd7aef0929ff07713255 Mon Sep 17 00:00:00 2001
From: Titouan Rigoudy <titouan.rigoudy@gmail.com>
Date: Mon, 4 Jan 2021 19:17:56 +0000
Subject: [PATCH] Add test for very large messages.

---
 src/proto/connection.rs | 24 +++++++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/src/proto/connection.rs b/src/proto/connection.rs
index 2cba8dd..0a7ee58 100644
--- a/src/proto/connection.rs
+++ b/src/proto/connection.rs
@@ -25,7 +25,7 @@ where
     pub fn new(stream: TcpStream) -> Self {
         Connection {
             stream,
-            read_buffer: BytesMut::with_capacity(4096),
+            read_buffer: BytesMut::new(),
             phantom_read: PhantomData,
             phantom_write: PhantomData,
         }
@@ -80,4 +80,26 @@ mod tests {
 
         server_task.await.unwrap();
     }
+
+    #[tokio::test]
+    async fn very_large_message() {
+        let listener = TcpListener::bind("localhost:0").await.unwrap();
+        let address = listener.local_addr().unwrap();
+
+        let server_task = tokio::spawn(async move {
+            let (stream, _peer_address) = listener.accept().await.unwrap();
+            let mut connection = Connection::<String, Vec<u32>>::new(stream);
+
+            assert_eq!(connection.read().await.unwrap(), "ping");
+            connection.write(&vec![0; 10 * 4096]).await.unwrap();
+        });
+
+        let stream = TcpStream::connect(address).await.unwrap();
+        let mut connection = Connection::<Vec<u32>, str>::new(stream);
+
+        connection.write("ping").await.unwrap();
+        assert_eq!(connection.read().await.unwrap(), vec![0; 10 * 4096]);
+
+        server_task.await.unwrap();
+    }
 }