From 8fc3b9522a4a56d881e94fc168279ac93c91e641 Mon Sep 17 00:00:00 2001
From: Mohammed Sazid Al Rashid <sazidozon@gmail.com>
Date: Thu, 21 Jan 2021 19:40:08 +0600
Subject: [PATCH] add tests for WsResponseBuilder

---
 actix-web-actors/tests/test_ws.rs | 122 +++++++++++++++++++++++++++---
 1 file changed, 111 insertions(+), 11 deletions(-)

diff --git a/actix-web-actors/tests/test_ws.rs b/actix-web-actors/tests/test_ws.rs
index 7fd59a4a7..d196c5470 100644
--- a/actix-web-actors/tests/test_ws.rs
+++ b/actix-web-actors/tests/test_ws.rs
@@ -1,4 +1,5 @@
 use actix::prelude::*;
+use actix_http::ws::Codec;
 use actix_web::{test, web, App, HttpRequest};
 use actix_web_actors::*;
 use bytes::Bytes;
@@ -26,16 +27,10 @@ impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for Ws {
     }
 }
 
-#[actix_rt::test]
-async fn test_simple() {
-    let mut srv = test::start(|| {
-        App::new().service(web::resource("/").to(
-            |req: HttpRequest, stream: web::Payload| async move {
-                ws::start(Ws, &req, stream)
-            },
-        ))
-    });
+const MAX_FRAME_SIZE: usize = 10_000;
+const DEFAULT_FRAME_SIZE: usize = 10;
 
+async fn common_test_code(mut srv: test::TestServer, frame_size: usize) {
     // client service
     let mut framed = srv.ws().await.unwrap();
     framed.send(ws::Message::Text("text".into())).await.unwrap();
@@ -43,12 +38,13 @@ async fn test_simple() {
     let item = framed.next().await.unwrap().unwrap();
     assert_eq!(item, ws::Frame::Text(Bytes::from_static(b"text")));
 
+    let bytes = Bytes::from(vec![0; frame_size]);
     framed
-        .send(ws::Message::Binary("text".into()))
+        .send(ws::Message::Binary(bytes.clone()))
         .await
         .unwrap();
     let item = framed.next().await.unwrap().unwrap();
-    assert_eq!(item, ws::Frame::Binary(Bytes::from_static(b"text")));
+    assert_eq!(item, ws::Frame::Binary(bytes));
 
     framed.send(ws::Message::Ping("text".into())).await.unwrap();
     let item = framed.next().await.unwrap().unwrap();
@@ -62,3 +58,107 @@ async fn test_simple() {
     let item = framed.next().await.unwrap().unwrap();
     assert_eq!(item, ws::Frame::Close(Some(ws::CloseCode::Normal.into())));
 }
+
+#[actix_rt::test]
+async fn test_builder() {
+    let srv = test::start(|| {
+        App::new().service(web::resource("/").to(
+            |req: HttpRequest, stream: web::Payload| async move {
+                ws::WsResponseBuilder::new(Ws, &req, stream).start()
+            },
+        ))
+    });
+
+    common_test_code(srv, DEFAULT_FRAME_SIZE).await;
+}
+
+#[actix_rt::test]
+async fn test_builder_with_frame_size() {
+    let srv = test::start(|| {
+        App::new().service(web::resource("/").to(
+            |req: HttpRequest, stream: web::Payload| async move {
+                ws::WsResponseBuilder::new(Ws, &req, stream)
+                    .frame_size(MAX_FRAME_SIZE)
+                    .start()
+            },
+        ))
+    });
+
+    common_test_code(srv, MAX_FRAME_SIZE).await;
+}
+
+#[actix_rt::test]
+async fn test_builder_with_codec() {
+    let srv = test::start(|| {
+        App::new().service(web::resource("/").to(
+            |req: HttpRequest, stream: web::Payload| async move {
+                ws::WsResponseBuilder::new(Ws, &req, stream)
+                    .codec(Codec::new())
+                    .start()
+            },
+        ))
+    });
+
+    common_test_code(srv, DEFAULT_FRAME_SIZE).await;
+}
+
+#[actix_rt::test]
+async fn test_builder_with_protocols() {
+    let srv = test::start(|| {
+        App::new().service(web::resource("/").to(
+            |req: HttpRequest, stream: web::Payload| async move {
+                ws::WsResponseBuilder::new(Ws, &req, stream)
+                    .protocols(&["A", "B"])
+                    .start()
+            },
+        ))
+    });
+
+    common_test_code(srv, DEFAULT_FRAME_SIZE).await;
+}
+
+#[actix_rt::test]
+async fn test_builder_full() {
+    let srv = test::start(|| {
+        App::new().service(web::resource("/").to(
+            |req: HttpRequest, stream: web::Payload| async move {
+                ws::WsResponseBuilder::new(Ws, &req, stream)
+                    .frame_size(MAX_FRAME_SIZE)
+                    .codec(Codec::new())
+                    .protocols(&["A", "B"])
+                    .start()
+            },
+        ))
+    });
+
+    common_test_code(srv, MAX_FRAME_SIZE).await;
+}
+
+#[actix_rt::test]
+async fn test_builder_with_codec_and_frame_size() {
+    let srv = test::start(|| {
+        App::new().service(web::resource("/").to(
+            |req: HttpRequest, stream: web::Payload| async move {
+                ws::WsResponseBuilder::new(Ws, &req, stream)
+                    .codec(Codec::new())
+                    .frame_size(MAX_FRAME_SIZE)
+                    .start()
+            },
+        ))
+    });
+
+    common_test_code(srv, DEFAULT_FRAME_SIZE).await;
+}
+
+#[actix_rt::test]
+async fn test_simple() {
+    let srv = test::start(|| {
+        App::new().service(web::resource("/").to(
+            |req: HttpRequest, stream: web::Payload| async move {
+                ws::start(Ws, &req, stream)
+            },
+        ))
+    });
+
+    common_test_code(srv, DEFAULT_FRAME_SIZE).await;
+}