diff --git a/CHANGES.md b/CHANGES.md index 5df0b6d6..39743be2 100644 --- a/CHANGES.md +++ b/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] diff --git a/actix-http-test/CHANGES.md b/actix-http-test/CHANGES.md index d6a2cdd9..0fac84a6 100644 --- a/actix-http-test/CHANGES.md +++ b/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. diff --git a/actix-http-test/src/lib.rs b/actix-http-test/src/lib.rs index 3749b78c..9a5069c4 100644 --- a/actix-http-test/src/lib.rs +++ b/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(); diff --git a/actix-web-actors/Cargo.toml b/actix-web-actors/Cargo.toml index 1f734582..e60d4301 100644 --- a/actix-web-actors/Cargo.toml +++ b/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 } diff --git a/actix-web-actors/tests/test_ws.rs b/actix-web-actors/tests/test_ws.rs index 912480ae..7999beed 100644 --- a/actix-web-actors/tests/test_ws.rs +++ b/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()))); +} diff --git a/src/test.rs b/src/test.rs index 7e8bf106..18bf89cd 100644 --- a/src/test.rs +++ b/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;