From f277b128b65c874b46ff897e4b770f8a713b88a4 Mon Sep 17 00:00:00 2001 From: fakeshadow <24548779@qq.com> Date: Thu, 13 May 2021 19:24:32 +0800 Subject: [PATCH] cleanup ws test (#2213) --- actix-http/tests/test_ws.rs | 182 +++++++++++++++--------------------- 1 file changed, 76 insertions(+), 106 deletions(-) diff --git a/actix-http/tests/test_ws.rs b/actix-http/tests/test_ws.rs index bf1ca938..b17d4211 100644 --- a/actix-http/tests/test_ws.rs +++ b/actix-http/tests/test_ws.rs @@ -1,193 +1,163 @@ -use std::cell::Cell; -use std::future::Future; -use std::marker::PhantomData; -use std::pin::Pin; -use std::sync::{Arc, Mutex}; -use std::task::{Context, Poll}; +use std::{ + cell::Cell, + task::{Context, Poll}, +}; 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_service::{fn_factory, Service}; -use actix_utils::future; use bytes::Bytes; +use futures_core::future::LocalBoxFuture; use futures_util::{SinkExt as _, StreamExt as _}; -use crate::ws::Dispatcher; +#[derive(Clone)] +struct WsService(Cell); -struct WsService(Arc, Cell)>>); - -impl WsService { +impl WsService { fn new() -> Self { - WsService(Arc::new(Mutex::new((PhantomData, Cell::new(false))))) + WsService(Cell::new(false)) } fn set_polled(&self) { - *self.0.lock().unwrap().1.get_mut() = true; + self.0.set(true); } fn was_polled(&self) -> bool { - self.0.lock().unwrap().1.get() + self.0.get() } } -impl Clone for WsService { - fn clone(&self) -> Self { - WsService(self.0.clone()) - } -} - -impl Service<(Request, Framed)> for WsService +impl Service<(Request, Framed)> for WsService where T: AsyncRead + AsyncWrite + Unpin + 'static, { type Response = (); type Error = Error; - type Future = Pin>>>; + type Future = LocalBoxFuture<'static, Result>; - fn poll_ready(&self, _ctx: &mut Context<'_>) -> Poll> { + fn poll_ready(&self, _: &mut Context<'_>) -> Poll> { self.set_polled(); Poll::Ready(Ok(())) } fn call(&self, (req, mut framed): (Request, Framed)) -> Self::Future { - let fut = async move { - let res = ws::handshake(req.head()).unwrap().message_body(()).unwrap(); + assert!(self.was_polled()); - framed - .send((res, body::BodySize::None).into()) - .await - .unwrap(); + Box::pin(async move { + let res = ws::handshake(req.head())?.message_body(())?; - Dispatcher::with(framed.replace_codec(ws::Codec::new()), service) - .await - .map_err(|_| panic!()) - }; + framed.send((res, BodySize::None).into()).await?; - 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 { +async fn service(msg: Frame) -> Result { let msg = match msg { - ws::Frame::Ping(msg) => ws::Message::Pong(msg), - ws::Frame::Text(text) => { - ws::Message::Text(String::from_utf8_lossy(&text).into_owned().into()) + Frame::Ping(msg) => Message::Pong(msg), + Frame::Text(text) => { + Message::Text(String::from_utf8_lossy(&text).into_owned().into()) } - ws::Frame::Binary(bin) => ws::Message::Binary(bin), - ws::Frame::Continuation(item) => ws::Message::Continuation(item), - ws::Frame::Close(reason) => ws::Message::Close(reason), - _ => panic!(), + Frame::Binary(bin) => Message::Binary(bin), + Frame::Continuation(item) => Message::Continuation(item), + Frame::Close(reason) => Message::Close(reason), + _ => return Err(Error::from(ws::ProtocolError::BadOpCode)), }; + Ok(msg) } #[actix_rt::test] async fn test_simple() { - let ws_service = WsService::new(); - let mut srv = test_server({ - let ws_service = ws_service.clone(); - move || { - let ws_service = ws_service.clone(); - HttpService::build() - .upgrade(fn_factory(move || future::ok::<_, ()>(ws_service.clone()))) - .finish(|_| future::ok::<_, ()>(Response::not_found())) - .tcp() - } + let mut srv = test_server(|| { + HttpService::build() + .upgrade(fn_factory(|| async { Ok::<_, ()>(WsService::new()) })) + .finish(|_| async { Ok::<_, ()>(Response::not_found()) }) + .tcp() }) .await; // client service let mut framed = srv.ws().await.unwrap(); - framed.send(ws::Message::Text("text".into())).await.unwrap(); - let (item, mut framed) = framed.into_future().await; - assert_eq!( - item.unwrap().unwrap(), - ws::Frame::Text(Bytes::from_static(b"text")) - ); + framed.send(Message::Text("text".into())).await.unwrap(); + + let item = framed.next().await.unwrap().unwrap(); + assert_eq!(item, 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 - .send(ws::Message::Binary("text".into())) + .send(Message::Continuation(Item::FirstText("text".into()))) .await .unwrap(); - let (item, mut framed) = framed.into_future().await; + let item = framed.next().await.unwrap().unwrap(); assert_eq!( - item.unwrap().unwrap(), - ws::Frame::Binary(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"))) + item, + Frame::Continuation(Item::FirstText(Bytes::from_static(b"text"))) ); assert!(framed - .send(ws::Message::Continuation(ws::Item::FirstText( - "text".into() - ))) + .send(Message::Continuation(Item::FirstText("text".into()))) .await .is_err()); assert!(framed - .send(ws::Message::Continuation(ws::Item::FirstBinary( - "text".into() - ))) + .send(Message::Continuation(Item::FirstBinary("text".into()))) .await .is_err()); framed - .send(ws::Message::Continuation(ws::Item::Continue("text".into()))) + .send(Message::Continuation(Item::Continue("text".into()))) .await .unwrap(); - let (item, mut framed) = framed.into_future().await; + let item = framed.next().await.unwrap().unwrap(); assert_eq!( - item.unwrap().unwrap(), - ws::Frame::Continuation(ws::Item::Continue(Bytes::from_static(b"text"))) + item, + Frame::Continuation(Item::Continue(Bytes::from_static(b"text"))) ); framed - .send(ws::Message::Continuation(ws::Item::Last("text".into()))) + .send(Message::Continuation(Item::Last("text".into()))) .await .unwrap(); - let (item, mut framed) = framed.into_future().await; + let item = framed.next().await.unwrap().unwrap(); assert_eq!( - item.unwrap().unwrap(), - ws::Frame::Continuation(ws::Item::Last(Bytes::from_static(b"text"))) + item, + Frame::Continuation(Item::Last(Bytes::from_static(b"text"))) ); assert!(framed - .send(ws::Message::Continuation(ws::Item::Continue("text".into()))) + .send(Message::Continuation(Item::Continue("text".into()))) .await .is_err()); assert!(framed - .send(ws::Message::Continuation(ws::Item::Last("text".into()))) + .send(Message::Continuation(Item::Last("text".into()))) .await .is_err()); framed - .send(ws::Message::Close(Some(ws::CloseCode::Normal.into()))) + .send(Message::Close(Some(CloseCode::Normal.into()))) .await .unwrap(); - let (item, _framed) = framed.into_future().await; - assert_eq!( - item.unwrap().unwrap(), - ws::Frame::Close(Some(ws::CloseCode::Normal.into())) - ); - - assert!(ws_service.was_polled()); + let item = framed.next().await.unwrap().unwrap(); + assert_eq!(item, Frame::Close(Some(CloseCode::Normal.into()))); }