From b0b207e3e80359dfb0829ce4536c69b46935d1d6 Mon Sep 17 00:00:00 2001
From: Rob Ede <robjtede@icloud.com>
Date: Thu, 9 Dec 2021 12:59:00 +0000
Subject: [PATCH] add take_complete_body tests

---
 actix-http/src/body/message_body.rs | 47 ++++++++++++++++++++++++++---
 examples/basic.rs                   |  4 +--
 2 files changed, 44 insertions(+), 7 deletions(-)

diff --git a/actix-http/src/body/message_body.rs b/actix-http/src/body/message_body.rs
index 47d68231b..e4020d2af 100644
--- a/actix-http/src/body/message_body.rs
+++ b/actix-http/src/body/message_body.rs
@@ -54,6 +54,9 @@ pub trait MessageBody {
     /// - A call to [`poll_next`] after calling [`take_complete_body`] should return `None` unless
     ///   the chunk is guaranteed to be empty.
     ///
+    /// The default implementation panics unconditionally, indicating a control flow bug in the
+    /// calling code.
+    ///
     /// # Panics
     /// With a correct implementation, panics if called without first checking [`is_complete_body`].
     ///
@@ -562,13 +565,49 @@ mod tests {
         assert_poll_next!(pl, Bytes::from("test"));
     }
 
-    #[actix_rt::test]
-    async fn big_string() {
-        let mut data = "HELLOWORLD".repeat(64 * 1024);
+    #[test]
+    fn take_string() {
+        let mut data = "test".repeat(2);
         let data_bytes = Bytes::from(data.clone());
-
         assert!(data.is_complete_body());
         assert_eq!(data.take_complete_body(), data_bytes);
+
+        let mut big_data = "test".repeat(64 * 1024);
+        let data_bytes = Bytes::from(big_data.clone());
+        assert!(big_data.is_complete_body());
+        assert_eq!(big_data.take_complete_body(), data_bytes);
+    }
+
+    #[test]
+    fn take_boxed_equivalence() {
+        let mut data = Bytes::from_static(b"test");
+        assert!(data.is_complete_body());
+        assert_eq!(data.take_complete_body(), b"test".as_ref());
+
+        let mut data = Box::new(Bytes::from_static(b"test"));
+        assert!(data.is_complete_body());
+        assert_eq!(data.take_complete_body(), b"test".as_ref());
+
+        let mut data = Box::pin(Bytes::from_static(b"test"));
+        assert!(data.is_complete_body());
+        assert_eq!(data.take_complete_body(), b"test".as_ref());
+    }
+
+    #[test]
+    fn take_policy() {
+        let mut data = Bytes::from_static(b"test");
+        // first call returns chunk
+        assert_eq!(data.take_complete_body(), b"test".as_ref());
+        // second call returns empty
+        assert_eq!(data.take_complete_body(), b"".as_ref());
+
+        let waker = futures_util::task::noop_waker();
+        let mut cx = Context::from_waker(&waker);
+        let mut data = Bytes::from_static(b"test");
+        // take returns whole chunk
+        assert_eq!(data.take_complete_body(), b"test".as_ref());
+        // subsequent poll_next returns None
+        assert_eq!(Pin::new(&mut data).poll_next(&mut cx), Poll::Ready(None));
     }
 
     // down-casting used to be done with a method on MessageBody trait
diff --git a/examples/basic.rs b/examples/basic.rs
index d10abb0a1..d29546129 100644
--- a/examples/basic.rs
+++ b/examples/basic.rs
@@ -11,11 +11,9 @@ async fn index_async(req: HttpRequest) -> &'static str {
     "Hello world!\r\n"
 }
 
-static A: [u8; 4096] = [b'0'; 4096];
-
 #[get("/")]
 async fn no_params() -> &'static str {
-    std::str::from_utf8(A.as_ref()).unwrap()
+    "Hello world!\r\n"
 }
 
 #[actix_web::main]