From ddd973b80839940e896b67c5fb2bcc6138c2654e Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Mon, 22 Feb 2021 18:33:47 +0600
Subject: [PATCH] framed state tests

---
 ntex/src/framed/dispatcher.rs |  3 +-
 ntex/src/framed/state.rs      | 78 +++++++++++++++++++++++++++++++++++
 ntex/src/testing.rs           |  5 +++
 3 files changed, 84 insertions(+), 2 deletions(-)

diff --git a/ntex/src/framed/dispatcher.rs b/ntex/src/framed/dispatcher.rs
index 922b76bc..de4110e6 100644
--- a/ntex/src/framed/dispatcher.rs
+++ b/ntex/src/framed/dispatcher.rs
@@ -706,7 +706,6 @@ mod tests {
 
     #[ntex_rt::test]
     async fn test_keepalive() {
-        env_logger::init();
         let (client, server) = Io::create();
         // do not allow to write to socket
         client.remote_buffer_cap(1024);
@@ -735,7 +734,7 @@ mod tests {
                 }
             }),
         );
-        crate::rt::spawn(disp.map(|_| ()));
+        crate::rt::spawn(disp.keepalive_timeout(0).keepalive_timeout(1).map(|_| ()));
 
         let state = state.disconnect_timeout(1);
 
diff --git a/ntex/src/framed/state.rs b/ntex/src/framed/state.rs
index 3c227dd7..792894fd 100644
--- a/ntex/src/framed/state.rs
+++ b/ntex/src/framed/state.rs
@@ -614,3 +614,81 @@ impl State {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use bytes::Bytes;
+
+    use crate::codec::BytesCodec;
+    use crate::testing::Io;
+
+    use super::*;
+
+    const BIN: &[u8] = b"GET /test HTTP/1\r\n\r\n";
+    const TEXT: &str = "GET /test HTTP/1\r\n\r\n";
+
+    #[ntex_rt::test]
+    async fn test_utils() {
+        let (client, mut server) = Io::create();
+        client.remote_buffer_cap(1024);
+        client.write(TEXT);
+
+        let state = State::new();
+        assert!(!state.is_read_buf_full());
+        assert!(!state.is_write_buf_full());
+
+        let msg = state.next(&mut server, &BytesCodec).await.unwrap().unwrap();
+        assert_eq!(msg, Bytes::from_static(BIN));
+
+        let res =
+            poll_fn(|cx| Poll::Ready(state.poll_next(&mut server, &BytesCodec, cx)))
+                .await;
+        assert!(res.is_pending());
+        client.write(TEXT);
+        let res =
+            poll_fn(|cx| Poll::Ready(state.poll_next(&mut server, &BytesCodec, cx)))
+                .await;
+        if let Poll::Ready(msg) = res {
+            assert_eq!(msg.unwrap().unwrap(), Bytes::from_static(BIN));
+        }
+
+        client.read_error(io::Error::new(io::ErrorKind::Other, "err"));
+        let msg = state.next(&mut server, &BytesCodec).await;
+        assert!(msg.is_err());
+        state.flags().contains(Flags::IO_ERR);
+        state.flags().contains(Flags::DSP_STOP);
+        state.remove_flags(Flags::IO_ERR | Flags::DSP_STOP);
+
+        client.read_error(io::Error::new(io::ErrorKind::Other, "err"));
+        let res =
+            poll_fn(|cx| Poll::Ready(state.poll_next(&mut server, &BytesCodec, cx)))
+                .await;
+        if let Poll::Ready(msg) = res {
+            assert!(msg.is_err());
+            state.flags().contains(Flags::IO_ERR);
+            state.flags().contains(Flags::DSP_STOP);
+            state.remove_flags(Flags::IO_ERR | Flags::DSP_STOP);
+        }
+
+        state
+            .send(&mut server, &BytesCodec, Bytes::from_static(b"test"))
+            .await
+            .unwrap();
+        let buf = client.read().await.unwrap();
+        assert_eq!(buf, Bytes::from_static(b"test"));
+
+        client.write_error(io::Error::new(io::ErrorKind::Other, "err"));
+        let res = state
+            .send(&mut server, &BytesCodec, Bytes::from_static(b"test"))
+            .await;
+        assert!(res.is_err());
+        state.flags().contains(Flags::IO_ERR);
+        state.flags().contains(Flags::DSP_STOP);
+        state.remove_flags(Flags::IO_ERR | Flags::DSP_STOP);
+
+        state.remove_flags(Flags::IO_ERR | Flags::DSP_STOP);
+        state.shutdown();
+        state.flags().contains(Flags::DSP_STOP);
+        state.flags().contains(Flags::IO_SHUTDOWN);
+    }
+}
diff --git a/ntex/src/testing.rs b/ntex/src/testing.rs
index 60003e23..9cc654ca 100644
--- a/ntex/src/testing.rs
+++ b/ntex/src/testing.rs
@@ -123,6 +123,11 @@ impl Io {
         self.remote.lock().unwrap().borrow_mut().read = IoState::Err(err);
     }
 
+    /// Set write error on remote side
+    pub fn write_error(&self, err: io::Error) {
+        self.local.lock().unwrap().borrow_mut().write = IoState::Err(err);
+    }
+
     /// Access read buffer.
     pub fn local_buffer<F, R>(&self, f: F) -> R
     where