Testing the TCP Server

Let's move on to testing our handle_connection function.

First, we need a TcpStream to work with. In an end-to-end or integration test, we might want to make a real TCP connection to test our code. One strategy for doing this is to start a listener on localhost port 0. Port 0 isn't a valid UNIX port, but it'll work for testing. The operating system will pick an open TCP port for us.

Instead, in this example we'll write a unit test for the connection handler, to check that the correct responses are returned for the respective inputs. To keep our unit test isolated and deterministic, we'll replace the TcpStream with a mock.

First, we'll change the signature of handle_connection to make it easier to test. handle_connection doesn't actually require an async_std::net::TcpStream; it requires any struct that implements async_std::io::Read, async_std::io::Write, and marker::Unpin. Changing the type signature to reflect this allows us to pass a mock for testing.

use async_std::io::{Read, Write};

async fn handle_connection(mut stream: impl Read + Write + Unpin) {

Next, let's build a mock TcpStream that implements these traits. First, let's implement the Read trait, with one method, poll_read. Our mock TcpStream will contain some data that is copied into the read buffer, and we'll return Poll::Ready to signify that the read is complete.

    use super::*;
    use futures::io::Error;
    use futures::task::{Context, Poll};

    use std::cmp::min;
    use std::pin::Pin;

    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 = min(self.read_data.len(), buf.len());
            buf[..size].copy_from_slice(&self.read_data[..size]);
            Poll::Ready(Ok(size))
        }
    }

Our implementation of Write is very similar, although we'll need to write three methods: poll_write, poll_flush, and poll_close. poll_write will copy any input data into the mock TcpStream, and return Poll::Ready when complete. No work needs to be done to flush or close the mock TcpStream, so poll_flush and poll_close can just return Poll::Ready.

    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(()))
        }
    }

Lastly, our mock will need to implement Unpin, signifying that its location in memory can safely be moved. For more information on pinning and the Unpin trait, see the section on pinning.

    impl Unpin for MockTcpStream {}

Now we're ready to test the handle_connection function. After setting up the MockTcpStream containing some initial data, we can run handle_connection using the attribute #[async_std::test], similarly to how we used #[async_std::main]. To ensure that handle_connection works as intended, we'll check that the correct data was written to the MockTcpStream based on its initial contents.

    use std::fs;

    #[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("hello.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()));
    }