From 7c6235cdc417a455aca911c8f80a5900a131f7b6 Mon Sep 17 00:00:00 2001
From: jamescarterbell <38409055+jamescarterbell@users.noreply.github.com>
Date: Tue, 16 Jul 2024 16:40:02 -0400
Subject: [PATCH] Add Server to TestServer (#383)

* Add Server to TestServer

* Version bump

---------

Co-authored-by: James Bell <jamesbell@microsoft.com>
---
 ntex-server/CHANGES.md      |  4 ++++
 ntex-server/Cargo.toml      |  2 +-
 ntex-server/src/net/test.rs | 43 ++++++++++++++++++++++++-------------
 3 files changed, 33 insertions(+), 16 deletions(-)

diff --git a/ntex-server/CHANGES.md b/ntex-server/CHANGES.md
index 3587ccd8..33ca163e 100644
--- a/ntex-server/CHANGES.md
+++ b/ntex-server/CHANGES.md
@@ -1,5 +1,9 @@
 # Changes
 
+## [2.3.0] - 2024-07-16
+
+* Add Server to TestServer
+
 ## [2.2.0] - 2024-07-16
 
 * Expose TetServer::stop() method
diff --git a/ntex-server/Cargo.toml b/ntex-server/Cargo.toml
index 4f876472..6f7874a5 100644
--- a/ntex-server/Cargo.toml
+++ b/ntex-server/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "ntex-server"
-version = "2.2.0"
+version = "2.3.0"
 authors = ["ntex contributors <team@ntex.rs>"]
 description = "Server for ntex framework"
 keywords = ["network", "framework", "async", "futures"]
diff --git a/ntex-server/src/net/test.rs b/ntex-server/src/net/test.rs
index e569b664..80441628 100644
--- a/ntex-server/src/net/test.rs
+++ b/ntex-server/src/net/test.rs
@@ -1,12 +1,12 @@
 //! Test server
-use std::{io, net, sync::mpsc, thread};
+use std::{io, net, thread};
 
 use ntex_net::{tcp_connect, Io};
 use ntex_rt::System;
 use ntex_service::ServiceFactory;
 use socket2::{Domain, SockAddr, Socket, Type};
 
-use super::ServerBuilder;
+use super::{Server, ServerBuilder};
 
 /// Start test server
 ///
@@ -44,29 +44,34 @@ where
     F: Fn() -> R + Send + Clone + 'static,
     R: ServiceFactory<Io> + 'static,
 {
-    let (tx, rx) = mpsc::channel();
-
+    let (tx, rx) = oneshot::channel();
     // run server in separate thread
     thread::spawn(move || {
         let sys = System::new("ntex-test-server");
         let tcp = net::TcpListener::bind("127.0.0.1:0").unwrap();
         let local_addr = tcp.local_addr().unwrap();
+        let system = sys.system();
 
-        tx.send((sys.system(), local_addr)).unwrap();
-        sys.run(|| {
-            ServerBuilder::new()
+        sys.run(move || {
+            let server = ServerBuilder::new()
                 .listen("test", tcp, move |_| factory())?
                 .set_tag("test", "TEST-SERVER")
                 .workers(1)
                 .disable_signals()
                 .run();
+            tx.send((system, local_addr, server))
+                .expect("Failed to send Server to TestServer");
             Ok(())
         })
     });
 
-    let (system, addr) = rx.recv().unwrap();
+    let (system, addr, server) = rx.recv().unwrap();
 
-    TestServer { addr, system }
+    TestServer {
+        addr,
+        server,
+        system,
+    }
 }
 
 /// Start new server with server builder
@@ -74,22 +79,24 @@ pub fn build_test_server<F>(factory: F) -> TestServer
 where
     F: FnOnce(ServerBuilder) -> ServerBuilder + Send + 'static,
 {
-    let (tx, rx) = mpsc::channel();
-
+    let (tx, rx) = oneshot::channel();
     // run server in separate thread
     thread::spawn(move || {
         let sys = System::new("ntex-test-server");
+        let system = sys.system();
 
-        tx.send(sys.system()).unwrap();
         sys.run(|| {
-            factory(super::build()).workers(1).disable_signals().run();
+            let server = factory(super::build()).workers(1).disable_signals().run();
+            tx.send((system, server))
+                .expect("Failed to send Server to TestServer");
             Ok(())
         })
     });
-    let system = rx.recv().unwrap();
+    let (system, server) = rx.recv().unwrap();
 
     TestServer {
         system,
+        server,
         addr: "127.0.0.1:0".parse().unwrap(),
     }
 }
@@ -99,6 +106,7 @@ where
 pub struct TestServer {
     addr: net::SocketAddr,
     system: System,
+    server: Server,
 }
 
 impl TestServer {
@@ -117,7 +125,7 @@ impl TestServer {
         tcp_connect(self.addr).await
     }
 
-    /// Stop http server
+    /// Stop http server by stopping the runtime.
     pub fn stop(&self) {
         self.system.stop();
     }
@@ -131,6 +139,11 @@ impl TestServer {
         let tcp = net::TcpListener::from(socket);
         tcp.local_addr().unwrap()
     }
+
+    /// Get access to the running Server
+    pub fn server(&self) -> Server {
+        self.server.clone()
+    }
 }
 
 impl Drop for TestServer {