diff --git a/src/net/server/tests.rs b/src/net/server/tests.rs
index 9cc784ac..6ac63da9 100644
--- a/src/net/server/tests.rs
+++ b/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()))
     }