added TestServer::ws_at_with_config

actix · Mar 25, 2021 · 72949a2 · 72949a2
1 parent 9704bed
commit 72949a2
Show file tree

Hide file tree

Showing 6 changed files with 84 additions and 9 deletions.
diff --git a/CHANGES.md b/CHANGES.md
@@ -1,6 +1,8 @@
 # Changes
 
 ## Unreleased - 2021-xx-xx
+* Added TestServer::ws_at_with_config method
+
 ### Fixed
 * Double ampersand in Logger format is escaped correctly. [#2067]
 

diff --git a/actix-http-test/CHANGES.md b/actix-http-test/CHANGES.md
@@ -1,7 +1,7 @@
 # Changes
 
 ## Unreleased - 2021-xx-xx
-
+* Added TestServer::ws_at_with_config method
 
 ## 3.0.0-beta.3 - 2021-03-09
 * No notable changes.

diff --git a/actix-http-test/src/lib.rs b/actix-http-test/src/lib.rs
@@ -241,14 +241,27 @@ impl TestServer {
         response.body().limit(10_485_760).await
     }
 
+    /// Connect to WebSocket server at given path and configure WebsocketsRequest using provided function.
+    pub async fn ws_at_with_config<ConfFunc>(
+        &mut self,
+        path: &str,
+        config_func: ConfFunc,
+    ) -> Result<Framed<impl AsyncRead + AsyncWrite, ws::Codec>, awc::error::WsClientError>
+    where
+        ConfFunc: FnOnce(ws::WebsocketsRequest) -> ws::WebsocketsRequest,
+    {
+        let url = self.url(path);
+        let request = self.client.ws(url);
+        let connect = config_func(request).connect();
+        connect.await.map(|(_, framed)| framed)
+    }
+
     /// Connect to WebSocket server at a given path.
     pub async fn ws_at(
         &mut self,
         path: &str,
     ) -> Result<Framed<impl AsyncRead + AsyncWrite, ws::Codec>, awc::error::WsClientError> {
-        let url = self.url(path);
-        let connect = self.client.ws(url).connect();
-        connect.await.map(|(_, framed)| framed)
+        self.ws_at_with_config(path, |wsr| wsr).await
     }
 
     /// Connect to a WebSocket server.

diff --git a/actix-web-actors/Cargo.toml b/actix-web-actors/Cargo.toml
@@ -29,5 +29,7 @@ tokio = { version = "1", features = ["sync"] }
 
 [dev-dependencies]
 actix-rt = "2.1"
+awc = { version = "3.0.0-beta.3", default-features = false }
 env_logger = "0.8"
 futures-util = { version = "0.3.7", default-features = false }
+http = "0.2.2"
diff --git a/actix-web-actors/tests/test_ws.rs b/actix-web-actors/tests/test_ws.rs
@@ -1,8 +1,9 @@
 use actix::prelude::*;
-use actix_web::{test, web, App, HttpRequest};
+use actix_web::{test, web, App, HttpRequest, HttpResponse};
 use actix_web_actors::*;
 use bytes::Bytes;
 use futures_util::{SinkExt, StreamExt};
+use http::StatusCode;
 
 struct Ws;
 
@@ -56,3 +57,47 @@ 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_with_credentials() {
+    let mut srv = test::start(|| {
+        App::new().service(web::resource("/").to(
+            |req: HttpRequest, stream: web::Payload| async move {
+                if req.headers().contains_key("Authorization") {
+                    ws::start(Ws, &req, stream)
+                } else {
+                    Ok(HttpResponse::new(StatusCode::UNAUTHORIZED))
+                }
+            },
+        ))
+    });
+
+    // client service without credentials
+    match srv.ws().await {
+        Ok(_) => panic!("WebSocket client without credentials should panic"),
+        Err(awc::error::WsClientError::InvalidResponseStatus(status)) => {
+            assert_eq!(status, StatusCode::UNAUTHORIZED)
+        }
+        Err(e) => panic!("Invalid error from WebSocket client: {}", e),
+    }
+
+    // client service with credentials
+    let client = srv.ws_at_with_config("/", |wsr| {
+        wsr.set_header("Authorization", "Bearer Something")
+    });
+
+    let mut framed = client.await.unwrap();
+
+    framed.send(ws::Message::Text("text".into())).await.unwrap();
+
+    let item = framed.next().await.unwrap().unwrap();
+    assert_eq!(item, ws::Frame::Text(Bytes::from_static(b"text")));
+
+    framed
+        .send(ws::Message::Close(Some(ws::CloseCode::Normal.into())))
+        .await
+        .unwrap();
+
+    let item = framed.next().await.unwrap().unwrap();
+    assert_eq!(item, ws::Frame::Close(Some(ws::CloseCode::Normal.into())));
+}
diff --git a/src/test.rs b/src/test.rs
@@ -16,7 +16,7 @@ use actix_router::{Path, ResourceDef, Url};
 use actix_rt::{time::sleep, System};
 use actix_service::{map_config, IntoService, IntoServiceFactory, Service, ServiceFactory};
 use awc::error::PayloadError;
-use awc::{Client, ClientRequest, ClientResponse, Connector};
+use awc::{ws::WebsocketsRequest, Client, ClientRequest, ClientResponse, Connector};
 use bytes::{Bytes, BytesMut};
 use futures_core::Stream;
 use futures_util::future::ok;
@@ -945,14 +945,27 @@ impl TestServer {
         response.body().limit(10_485_760).await
     }
 
+    /// Connect to WebSocket server at given path and configure WebsocketsRequest using provided function.
+    pub async fn ws_at_with_config<ConfFunc>(
+        &mut self,
+        path: &str,
+        config_func: ConfFunc,
+    ) -> Result<Framed<impl AsyncRead + AsyncWrite, ws::Codec>, awc::error::WsClientError>
+    where
+        ConfFunc: FnOnce(WebsocketsRequest) -> WebsocketsRequest,
+    {
+        let url = self.url(path);
+        let request = self.client.ws(url);
+        let connect = config_func(request).connect();
+        connect.await.map(|(_, framed)| framed)
+    }
+
     /// Connect to WebSocket server at a given path.
     pub async fn ws_at(
         &mut self,
         path: &str,
     ) -> Result<Framed<impl AsyncRead + AsyncWrite, ws::Codec>, awc::error::WsClientError> {
-        let url = self.url(path);
-        let connect = self.client.ws(url).connect();
-        connect.await.map(|(_, framed)| framed)
+        self.ws_at_with_config(path, |wsr| wsr).await
     }
 
     /// Connect to a WebSocket server.