From 080b81eee9945090bdf9574a7cff06d8b2fdf1e7 Mon Sep 17 00:00:00 2001
From: Ximon Eighteen <3304436+ximon18@users.noreply.github.com>
Date: Mon, 22 Apr 2024 13:47:50 +0200
Subject: [PATCH] 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 | 22 ++++++++++++++++------
 1 file changed, 16 insertions(+), 6 deletions(-)

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