From d5776cd81b01186cbd5d544935cf68b3eed0d0f3 Mon Sep 17 00:00:00 2001
From: Titouan Rigoudy <titouan.rigoudy@gmail.com>
Date: Sat, 23 Jan 2021 13:17:19 +0100
Subject: [PATCH] Rework FakeServer initialization.

---
 src/proto/server/client.rs  | 13 +++---
 src/proto/server/testing.rs | 87 +++++++++++++++++++------------------
 2 files changed, 52 insertions(+), 48 deletions(-)

diff --git a/src/proto/server/client.rs b/src/proto/server/client.rs
index d9f2bf4..8d296eb 100644
--- a/src/proto/server/client.rs
+++ b/src/proto/server/client.rs
@@ -179,7 +179,7 @@ mod tests {
   use tokio::net;
   use tokio::sync::mpsc;
 
-  use crate::proto::server::testing::{fake_server, UserStatusMap};
+  use crate::proto::server::testing::{ServerBuilder, UserStatusMap};
   use crate::proto::server::{
     Credentials, ServerRequest, ServerResponse, UserStatusRequest,
     UserStatusResponse,
@@ -209,7 +209,7 @@ mod tests {
   async fn login() {
     init();
 
-    let (server, handle) = fake_server().await.unwrap();
+    let (server, handle) = ServerBuilder::default().bind().await.unwrap();
     let server_task = tokio::spawn(server.serve());
 
     let stream = net::TcpStream::connect(handle.address()).await.unwrap();
@@ -237,8 +237,11 @@ mod tests {
     let mut user_status_map = UserStatusMap::default();
     user_status_map.insert(response.clone());
 
-    let (server, handle) = fake_server().await.unwrap();
-    let server = server.with_user_status_map(user_status_map);
+    let (server, handle) = ServerBuilder::default()
+      .with_user_status_map(user_status_map)
+      .bind()
+      .await
+      .unwrap();
     let server_task = tokio::spawn(server.serve());
 
     let stream = net::TcpStream::connect(handle.address()).await.unwrap();
@@ -269,7 +272,7 @@ mod tests {
   async fn stream_closed() {
     init();
 
-    let (server, handle) = fake_server().await.unwrap();
+    let (server, handle) = ServerBuilder::default().bind().await.unwrap();
     let server_task = tokio::spawn(server.serve());
 
     let stream = net::TcpStream::connect(handle.address()).await.unwrap();
diff --git a/src/proto/server/testing.rs b/src/proto/server/testing.rs
index f1aa5fb..6967aa2 100644
--- a/src/proto/server/testing.rs
+++ b/src/proto/server/testing.rs
@@ -137,13 +137,46 @@ impl GracefulHandler {
   }
 }
 
-// TODO: Rework interface to look like:
-//
-// let (server, handle) = FakeServerBuilder::new()
-//   .with_shutdown(shutdown_rx)
-//   .with_user_status_map(map)
-//   .bind().await.unwrap();
-//
+/// A builder for FakeServer instances.
+#[derive(Default)]
+pub struct ServerBuilder {
+  user_status_map: Option<Arc<Mutex<UserStatusMap>>>,
+}
+
+impl ServerBuilder {
+  /// Sets the UserStatusMap which the server will use to respond to
+  /// UserStatusRequest messages.
+  pub fn with_user_status_map(mut self, map: UserStatusMap) -> Self {
+    self.user_status_map = Some(Arc::new(Mutex::new(map)));
+    self
+  }
+
+  /// Binds to a localhost port, then returns a server and its handle.
+  pub async fn bind(self) -> io::Result<(FakeServer, FakeServerHandle)> {
+    let listener = TcpListener::bind("localhost:0").await?;
+    let address = listener.local_addr()?;
+
+    let user_status_map = match self.user_status_map {
+      Some(user_status_map) => user_status_map,
+      None => Arc::new(Mutex::new(UserStatusMap::default())),
+    };
+
+    let (shutdown_tx, shutdown_rx) = watch::channel(());
+
+    Ok((
+      FakeServer {
+        listener,
+        shutdown_rx,
+        user_status_map,
+      },
+      FakeServerHandle {
+        shutdown_tx,
+        address,
+      },
+    ))
+  }
+}
+
 /// A fake server for connecting to in tests.
 pub struct FakeServer {
   listener: TcpListener,
@@ -172,24 +205,6 @@ impl FakeServerHandle {
 }
 
 impl FakeServer {
-  /// Creates a new fake server and binds it to a port on localhost.
-  ///
-  /// The returned server will stop serving when `shutdown_rx` is notified, or
-  /// when its sender is dropped.
-  pub async fn bind(shutdown_rx: watch::Receiver<()>) -> io::Result<Self> {
-    let listener = TcpListener::bind("localhost:0").await?;
-    Ok(FakeServer {
-      listener,
-      shutdown_rx,
-      user_status_map: Arc::new(Mutex::new(UserStatusMap::default())),
-    })
-  }
-
-  pub fn with_user_status_map(mut self, map: UserStatusMap) -> Self {
-    self.user_status_map = Arc::new(Mutex::new(map));
-    self
-  }
-
   /// Returns the address to which this server is bound.
   /// This is always localhost and a random port chosen by the OS.
   pub fn address(&self) -> io::Result<SocketAddr> {
@@ -239,26 +254,11 @@ impl FakeServer {
   }
 }
 
-pub async fn fake_server() -> io::Result<(FakeServer, FakeServerHandle)> {
-  let (shutdown_tx, shutdown_rx) = watch::channel(());
-  let server = FakeServer::bind(shutdown_rx).await?;
-
-  let address = server.address()?;
-
-  Ok((
-    server,
-    FakeServerHandle {
-      shutdown_tx,
-      address,
-    },
-  ))
-}
-
 #[cfg(test)]
 mod tests {
   use tokio::net::TcpStream;
 
-  use super::fake_server;
+  use super::ServerBuilder;
 
   // Enable capturing logs in tests.
   fn init() {
@@ -269,15 +269,16 @@ mod tests {
   async fn new_binds_to_localhost() {
     init();
 
-    let (server, _handle) = fake_server().await.unwrap();
+    let (server, handle) = ServerBuilder::default().bind().await.unwrap();
     assert!(server.address().unwrap().ip().is_loopback());
+    assert_eq!(server.address().unwrap(), handle.address());
   }
 
   #[tokio::test]
   async fn accepts_incoming_connections() {
     init();
 
-    let (server, handle) = fake_server().await.unwrap();
+    let (server, handle) = ServerBuilder::default().bind().await.unwrap();
     let server_task = tokio::spawn(server.serve());
 
     // The connection succeeds.