From b68e0d3408450f97832e8a1fcbc673cf3ecbd15f Mon Sep 17 00:00:00 2001
From: Martin Pitt <martin@piware.de>
Date: Fri, 16 Sep 2022 09:39:05 +0200
Subject: [PATCH] async-http: Unit test for handle_connection()

---
 async-http/src/main.rs | 72 ++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 70 insertions(+), 2 deletions(-)

diff --git a/async-http/src/main.rs b/async-http/src/main.rs
index c288408..1a8dddb 100644
--- a/async-http/src/main.rs
+++ b/async-http/src/main.rs
@@ -2,7 +2,8 @@ use std::fs;
 use std::time::Duration;
 
 use async_std::prelude::*;
-use async_std::net::{ TcpListener, TcpStream };
+use async_std::io::{ Read, Write };
+use async_std::net::{ TcpListener };
 use async_std::task;
 use futures::stream::StreamExt;
 
@@ -19,7 +20,7 @@ async fn main() {
     }).await;
 }
 
-async fn handle_connection(mut stream: TcpStream) {
+async fn handle_connection(mut stream: impl Read + Write + Unpin) {
     // Read the first 1024 bytes of data from the stream
     let mut buffer = [0; 1024];
     assert!(stream.read(&mut buffer).await.unwrap() > 0);
@@ -45,3 +46,70 @@ async fn handle_connection(mut stream: TcpStream) {
     stream.write_all(response.as_bytes()).await.unwrap();
     stream.flush().await.unwrap();
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use std::cmp;
+    use std::pin::Pin;
+
+    use futures::io::Error;
+    use futures::task::{Context, Poll};
+
+    struct MockTcpStream {
+        read_data: Vec<u8>,
+        write_data: Vec<u8>,
+    }
+
+    impl Read for MockTcpStream {
+        fn poll_read(
+            self: Pin<&mut Self>,
+            _: &mut Context,
+            buf: &mut [u8],
+        ) -> Poll<Result<usize, Error>> {
+            let size: usize = cmp::min(self.read_data.len(), buf.len());
+            buf[..size].copy_from_slice(&self.read_data[..size]);
+            Poll::Ready(Ok(size))
+        }
+    }
+
+    impl Write for MockTcpStream {
+        fn poll_write(
+            mut self: Pin<&mut Self>,
+            _: &mut Context,
+            buf: &[u8],
+        ) -> Poll<Result<usize, Error>> {
+            self.write_data = Vec::from(buf);
+
+            Poll::Ready(Ok(buf.len()))
+        }
+
+        fn poll_flush(self: Pin<&mut Self>, _: &mut Context) -> Poll<Result<(), Error>> {
+            Poll::Ready(Ok(()))
+        }
+
+        fn poll_close(self: Pin<&mut Self>, _: &mut Context) -> Poll<Result<(), Error>> {
+            Poll::Ready(Ok(()))
+        }
+    }
+
+    impl Unpin for MockTcpStream {}
+
+    #[async_std::test]
+    async fn test_handle_connection() {
+        let input_bytes = b"GET / HTTP/1.1\r\n";
+        let mut contents = vec![0u8; 1024];
+        contents[..input_bytes.len()].clone_from_slice(input_bytes);
+        let mut stream = MockTcpStream {
+            read_data: contents,
+            write_data: Vec::new(),
+        };
+
+        handle_connection(&mut stream).await;
+
+        let expected_contents = fs::read_to_string("index.html").unwrap();
+        let expected_response = format!("HTTP/1.1 200 OK\r\n\r\n{}", expected_contents);
+        assert!(stream.write_data.starts_with(expected_response.as_bytes()));
+    }
+}
-- 
2.39.5