actix/actix-web
1
0
Fork 0
mirror of https://github.com/fafhrd91/actix-web synced 2024-11-27 17:52:56 +01:00
Code Issues Releases Wiki Activity

cleanup ws test (#2213)

Browse Source
This commit is contained in:
fakeshadow 2021-05-13 19:24:32 +08:00 committed by GitHub
parent 4903950b22
commit f277b128b6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 76 additions and 106 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

178
actix-http/tests/test_ws.rs
View File

@ -1,193 +1,163 @@
use std::cell::Cell; use std::{
use std::future::Future; cell::Cell,
use std::marker::PhantomData; task::{Context, Poll},
use std::pin::Pin; };
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use actix_codec::{AsyncRead, AsyncWrite, Framed}; use actix_codec::{AsyncRead, AsyncWrite, Framed};
use actix_http::{body, h1, ws, Error, HttpService, Request, Response}; use actix_http::{
body::BodySize,
h1,
ws::{self, CloseCode, Frame, Item, Message},
Error, HttpService, Request, Response,
};
use actix_http_test::test_server; use actix_http_test::test_server;
use actix_service::{fn_factory, Service}; use actix_service::{fn_factory, Service};
use actix_utils::future;
use bytes::Bytes; use bytes::Bytes;
use futures_core::future::LocalBoxFuture;
use futures_util::{SinkExt as _, StreamExt as _}; use futures_util::{SinkExt as _, StreamExt as _};
use crate::ws::Dispatcher; #[derive(Clone)]
struct WsService(Cell<bool>);
struct WsService<T>(Arc<Mutex<(PhantomData<T>, Cell<bool>)>>); impl WsService {
impl<T> WsService<T> {
fn new() -> Self { fn new() -> Self {
WsService(Arc::new(Mutex::new((PhantomData, Cell::new(false))))) WsService(Cell::new(false))
} }
fn set_polled(&self) { fn set_polled(&self) {
*self.0.lock().unwrap().1.get_mut() = true; self.0.set(true);
} }
fn was_polled(&self) -> bool { fn was_polled(&self) -> bool {
self.0.lock().unwrap().1.get() self.0.get()
} }
} }
impl<T> Clone for WsService<T> { impl<T> Service<(Request, Framed<T, h1::Codec>)> for WsService
fn clone(&self) -> Self {
WsService(self.0.clone())
}
}
impl<T> Service<(Request, Framed<T, h1::Codec>)> for WsService<T>
where where
T: AsyncRead + AsyncWrite + Unpin + 'static, T: AsyncRead + AsyncWrite + Unpin + 'static,
{ {
type Response = (); type Response = ();
type Error = Error; type Error = Error;
type Future = Pin<Box<dyn Future<Output = Result<(), Error>>>>; type Future = LocalBoxFuture<'static, Result<Self::Response, Self::Error>>;
fn poll_ready(&self, _ctx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { fn poll_ready(&self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.set_polled(); self.set_polled();
Poll::Ready(Ok(())) Poll::Ready(Ok(()))
} }
fn call(&self, (req, mut framed): (Request, Framed<T, h1::Codec>)) -> Self::Future { fn call(&self, (req, mut framed): (Request, Framed<T, h1::Codec>)) -> Self::Future {
let fut = async move { assert!(self.was_polled());
let res = ws::handshake(req.head()).unwrap().message_body(()).unwrap();
framed Box::pin(async move {
.send((res, body::BodySize::None).into()) let res = ws::handshake(req.head())?.message_body(())?;
.await
.unwrap();
Dispatcher::with(framed.replace_codec(ws::Codec::new()), service) framed.send((res, BodySize::None).into()).await?;
.await
.map_err(|_| panic!())
};
Box::pin(fut) let framed = framed.replace_codec(ws::Codec::new());
ws::Dispatcher::with(framed, service).await?;
Ok(())
})
} }
} }
async fn service(msg: ws::Frame) -> Result<ws::Message, Error> { async fn service(msg: Frame) -> Result<Message, Error> {
let msg = match msg { let msg = match msg {
ws::Frame::Ping(msg) => ws::Message::Pong(msg), Frame::Ping(msg) => Message::Pong(msg),
ws::Frame::Text(text) => { Frame::Text(text) => {
ws::Message::Text(String::from_utf8_lossy(&text).into_owned().into()) Message::Text(String::from_utf8_lossy(&text).into_owned().into())
} }
ws::Frame::Binary(bin) => ws::Message::Binary(bin), Frame::Binary(bin) => Message::Binary(bin),
ws::Frame::Continuation(item) => ws::Message::Continuation(item), Frame::Continuation(item) => Message::Continuation(item),
ws::Frame::Close(reason) => ws::Message::Close(reason), Frame::Close(reason) => Message::Close(reason),
_ => panic!(), _ => return Err(Error::from(ws::ProtocolError::BadOpCode)),
}; };
Ok(msg) Ok(msg)
} }
#[actix_rt::test] #[actix_rt::test]
async fn test_simple() { async fn test_simple() {
let ws_service = WsService::new(); let mut srv = test_server(|| {
let mut srv = test_server({
let ws_service = ws_service.clone();
move || {
let ws_service = ws_service.clone();
HttpService::build() HttpService::build()
.upgrade(fn_factory(move || future::ok::<_, ()>(ws_service.clone()))) .upgrade(fn_factory(|| async { Ok::<_, ()>(WsService::new()) }))
.finish(|_| future::ok::<_, ()>(Response::not_found())) .finish(|_| async { Ok::<_, ()>(Response::not_found()) })
.tcp() .tcp()
}
}) })
.await; .await;
// client service // client service
let mut framed = srv.ws().await.unwrap(); let mut framed = srv.ws().await.unwrap();
framed.send(ws::Message::Text("text".into())).await.unwrap(); framed.send(Message::Text("text".into())).await.unwrap();
let (item, mut framed) = framed.into_future().await;
assert_eq!( let item = framed.next().await.unwrap().unwrap();
item.unwrap().unwrap(), assert_eq!(item, Frame::Text(Bytes::from_static(b"text")));
ws::Frame::Text(Bytes::from_static(b"text"))
); framed.send(Message::Binary("text".into())).await.unwrap();
let item = framed.next().await.unwrap().unwrap();
assert_eq!(item, Frame::Binary(Bytes::from_static(&b"text"[..])));
framed.send(Message::Ping("text".into())).await.unwrap();
let item = framed.next().await.unwrap().unwrap();
assert_eq!(item, Frame::Pong("text".to_string().into()));
framed framed
.send(ws::Message::Binary("text".into())) .send(Message::Continuation(Item::FirstText("text".into())))
.await .await
.unwrap(); .unwrap();
let (item, mut framed) = framed.into_future().await; let item = framed.next().await.unwrap().unwrap();
assert_eq!( assert_eq!(
item.unwrap().unwrap(), item,
ws::Frame::Binary(Bytes::from_static(&b"text"[..])) Frame::Continuation(Item::FirstText(Bytes::from_static(b"text")))
);
framed.send(ws::Message::Ping("text".into())).await.unwrap();
let (item, mut framed) = framed.into_future().await;
assert_eq!(
item.unwrap().unwrap(),
ws::Frame::Pong("text".to_string().into())
);
framed
.send(ws::Message::Continuation(ws::Item::FirstText(
"text".into(),
)))
.await
.unwrap();
let (item, mut framed) = framed.into_future().await;
assert_eq!(
item.unwrap().unwrap(),
ws::Frame::Continuation(ws::Item::FirstText(Bytes::from_static(b"text")))
); );
assert!(framed assert!(framed
.send(ws::Message::Continuation(ws::Item::FirstText( .send(Message::Continuation(Item::FirstText("text".into())))
"text".into()
)))
.await .await
.is_err()); .is_err());
assert!(framed assert!(framed
.send(ws::Message::Continuation(ws::Item::FirstBinary( .send(Message::Continuation(Item::FirstBinary("text".into())))
"text".into()
)))
.await .await
.is_err()); .is_err());
framed framed
.send(ws::Message::Continuation(ws::Item::Continue("text".into()))) .send(Message::Continuation(Item::Continue("text".into())))
.await .await
.unwrap(); .unwrap();
let (item, mut framed) = framed.into_future().await; let item = framed.next().await.unwrap().unwrap();
assert_eq!( assert_eq!(
item.unwrap().unwrap(), item,
ws::Frame::Continuation(ws::Item::Continue(Bytes::from_static(b"text"))) Frame::Continuation(Item::Continue(Bytes::from_static(b"text")))
); );
framed framed
.send(ws::Message::Continuation(ws::Item::Last("text".into()))) .send(Message::Continuation(Item::Last("text".into())))
.await .await
.unwrap(); .unwrap();
let (item, mut framed) = framed.into_future().await; let item = framed.next().await.unwrap().unwrap();
assert_eq!( assert_eq!(
item.unwrap().unwrap(), item,
ws::Frame::Continuation(ws::Item::Last(Bytes::from_static(b"text"))) Frame::Continuation(Item::Last(Bytes::from_static(b"text")))
); );
assert!(framed assert!(framed
.send(ws::Message::Continuation(ws::Item::Continue("text".into()))) .send(Message::Continuation(Item::Continue("text".into())))
.await .await
.is_err()); .is_err());
assert!(framed assert!(framed
.send(ws::Message::Continuation(ws::Item::Last("text".into()))) .send(Message::Continuation(Item::Last("text".into())))
.await .await
.is_err()); .is_err());
framed framed
.send(ws::Message::Close(Some(ws::CloseCode::Normal.into()))) .send(Message::Close(Some(CloseCode::Normal.into())))
.await .await
.unwrap(); .unwrap();
let (item, _framed) = framed.into_future().await; let item = framed.next().await.unwrap().unwrap();
assert_eq!( assert_eq!(item, Frame::Close(Some(CloseCode::Normal.into())));
item.unwrap().unwrap(),
ws::Frame::Close(Some(ws::CloseCode::Normal.into()))
);
assert!(ws_service.was_polled());
} }