Explicitly terminate mock client connections once they have received all of their responses, rather than leave them connected till the end of the test to shutdown when the server shuts down. (#295)

src/net/server/tests.rs

@@ -41,6 +41,8 @@ struct MockStream {
 
     /// The rate at which messages should be made available to the server.
     new_message_every: Duration,
+
+    pending_responses: usize,
 }
 
 impl MockStream {
@@ -48,10 +50,12 @@ impl MockStream {
         messages_to_read: VecDeque<Vec<u8>>,
         new_message_every: Duration,
     ) -> Self {
+        let pending_responses = messages_to_read.len();
         Self {
             last_ready: Mutex::new(Option::None),
             messages_to_read: Mutex::new(messages_to_read),
             new_message_every,
+            pending_responses,
         }
     }
 }
@@ -80,11 +84,13 @@ impl AsyncRead for MockStream {
                         last_ready.replace(Instant::now());
                         return Poll::Ready(Ok(()));
                     } else {
-                        // End of stream
-                        /*return Poll::Ready(Err(io::Error::new(
-                            io::ErrorKind::ConnectionAborted,
-                            "mock connection disconnect",
-                        )));*/
+                        // Disconnect once we've sent all of the requests AND received all of the responses.
+                        if self.pending_responses == 0 {
+                            return Poll::Ready(Err(io::Error::new(
+                                io::ErrorKind::ConnectionAborted,
+                                "mock connection disconnect",
+                            )));
+                        }
                     }
                 }
                 _ => {
@@ -109,10 +115,14 @@ impl AsyncRead for MockStream {
 
 impl AsyncWrite for MockStream {
     fn poll_write(
-        self: Pin<&mut Self>,
+        mut self: Pin<&mut Self>,
         _cx: &mut Context<'_>,
         buf: &[u8],
     ) -> Poll<Result<usize, io::Error>> {
+        // Assume a single write is an entire response.
+        if self.pending_responses > 0 {
+            self.pending_responses -= 1;
+        }
         Poll::Ready(Ok(buf.len()))
     }