added TestServer::client_headers (#2097)

Co-authored-by: fakeshadow <24548779@qq.com>
Co-authored-by: Rob Ede <robjtede@icloud.com>

CHANGES.md

@@ -1,6 +1,9 @@
 # Changes
 
 ## Unreleased - 2021-xx-xx
+### Added
+* Added `TestServer::client_headers` method. [#2097]
+
 ### Fixed
 * Double ampersand in Logger format is escaped correctly. [#2067]
 

actix-http-test/CHANGES.md

@@ -1,7 +1,8 @@
 # Changes
 
 ## Unreleased - 2021-xx-xx
-
+### Added
+* Added `TestServer::client_headers` method. [#2097]
 
 ## 3.0.0-beta.3 - 2021-03-09
 * No notable changes.

actix-http-test/src/lib.rs

@@ -13,7 +13,9 @@ use std::{net, thread, time};
 use actix_codec::{AsyncRead, AsyncWrite, Framed};
 use actix_rt::{net::TcpStream, System};
 use actix_server::{Server, ServiceFactory};
-use awc::{error::PayloadError, ws, Client, ClientRequest, ClientResponse, Connector};
+use awc::{
+    error::PayloadError, http::HeaderMap, ws, Client, ClientRequest, ClientResponse, Connector,
+};
 use bytes::Bytes;
 use futures_core::stream::Stream;
 use http::Method;
@@ -258,6 +260,14 @@ impl TestServer {
         self.ws_at("/").await
     }
 
+    /// Get default HeaderMap of Client.
+    ///
+    /// Returns Some(&mut HeaderMap) when Client object is unique
+    /// (No other clone of client exists at the same time).
+    pub fn client_headers(&mut self) -> Option<&mut HeaderMap> {
+        self.client.headers()
+    }
+
     /// Stop HTTP server
     fn stop(&mut self) {
         self.system.stop();

actix-web-actors/Cargo.toml

@@ -29,5 +29,6 @@ tokio = { version = "1", features = ["sync"] }
 
 [dev-dependencies]
 actix-rt = "2.2"
+awc = { version = "3.0.0-beta.3", default-features = false }
 env_logger = "0.8"
 futures-util = { version = "0.3.7", default-features = false }

actix-web-actors/tests/test_ws.rs

@@ -1,5 +1,8 @@
 use actix::prelude::*;
-use actix_web::{test, web, App, HttpRequest};
+use actix_web::{
+    http::{header, StatusCode},
+    test, web, App, HttpRequest, HttpResponse,
+};
 use actix_web_actors::*;
 use bytes::Bytes;
 use futures_util::{SinkExt, StreamExt};
@@ -56,3 +59,51 @@ 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),
+    }
+
+    let headers = srv.client_headers().unwrap();
+    headers.insert(
+        header::AUTHORIZATION,
+        header::HeaderValue::from_static("Bearer Something"),
+    );
+
+    // client service with credentials
+    let client = srv.ws();
+
+    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())));
+}

src/test.rs

@@ -8,7 +8,7 @@ use std::{fmt, net, thread, time};
 use actix_codec::{AsyncRead, AsyncWrite, Framed};
 #[cfg(feature = "cookies")]
 use actix_http::cookie::Cookie;
-use actix_http::http::header::{ContentType, IntoHeaderPair};
+use actix_http::http::header::{ContentType, HeaderMap, IntoHeaderPair};
 use actix_http::http::{Method, StatusCode, Uri, Version};
 use actix_http::test::TestRequest as HttpTestRequest;
 use actix_http::{ws, Extensions, HttpService, Request};
@@ -962,6 +962,14 @@ impl TestServer {
         self.ws_at("/").await
     }
 
+    /// Get default HeaderMap of Client.
+    ///
+    /// Returns Some(&mut HeaderMap) when Client object is unique
+    /// (No other clone of client exists at the same time).
+    pub fn client_headers(&mut self) -> Option<&mut HeaderMap> {
+        self.client.headers()
+    }
+
     /// Gracefully stop HTTP server
     pub async fn stop(self) {
         self.server.stop(true).await;