From 7882f545e5ff3efaadb2095ca2dfcbd59135855d Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Wed, 25 Dec 2019 12:10:28 +0400
Subject: [PATCH] Allow to gracefully stop test server via TestServer::stop()

---
 CHANGES.md  |  2 ++
 src/data.rs | 46 +++++++++++++++++++++++++++++++++++++++++++++-
 src/test.rs | 21 ++++++++++++---------
 3 files changed, 59 insertions(+), 10 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index 0291f75a..31ebee53 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -6,6 +6,8 @@
 
 * Rename `HttpServer::start()` to `HttpServer::run()`
 
+* Allow to gracefully stop test server via `TestServer::stop()`
+
 
 ## [2.0.0-rc] - 2019-12-20
 
diff --git a/src/data.rs b/src/data.rs
index 2867d68f..c3641894 100644
--- a/src/data.rs
+++ b/src/data.rs
@@ -136,10 +136,11 @@ impl<T: 'static> DataFactory for Data<T> {
 #[cfg(test)]
 mod tests {
     use actix_service::Service;
+    use std::sync::atomic::{AtomicUsize, Ordering};
 
     use super::*;
     use crate::http::StatusCode;
-    use crate::test::{init_service, TestRequest};
+    use crate::test::{self, init_service, TestRequest};
     use crate::{web, App, HttpResponse};
 
     #[actix_rt::test]
@@ -234,4 +235,47 @@ mod tests {
         let resp = srv.call(req).await.unwrap();
         assert_eq!(resp.status(), StatusCode::OK);
     }
+
+    #[actix_rt::test]
+    async fn test_data_drop() {
+        struct TestData(Arc<AtomicUsize>);
+
+        impl TestData {
+            fn new(inner: Arc<AtomicUsize>) -> Self {
+                let _ = inner.fetch_add(1, Ordering::SeqCst);
+                Self(inner)
+            }
+        }
+
+        impl Clone for TestData {
+            fn clone(&self) -> Self {
+                let inner = self.0.clone();
+                let _ = inner.fetch_add(1, Ordering::SeqCst);
+                Self(inner)
+            }
+        }
+
+        impl Drop for TestData {
+            fn drop(&mut self) {
+                let _ = self.0.fetch_sub(1, Ordering::SeqCst);
+            }
+        }
+
+        let num = Arc::new(AtomicUsize::new(0));
+        let data = TestData::new(num.clone());
+        assert_eq!(num.load(Ordering::SeqCst), 1);
+
+        let srv = test::start(move || {
+            let data = data.clone();
+
+            App::new()
+                .data(data)
+                .service(web::resource("/").to(|_data: Data<TestData>| async { "ok" }))
+        });
+
+        assert!(srv.get("/").send().await.unwrap().status().is_success());
+        srv.stop().await;
+
+        assert_eq!(num.load(Ordering::SeqCst), 0);
+    }
 }
diff --git a/src/test.rs b/src/test.rs
index 912e9b47..9c06b733 100644
--- a/src/test.rs
+++ b/src/test.rs
@@ -11,8 +11,7 @@ use actix_http::http::{Error as HttpError, Method, StatusCode, Uri, Version};
 use actix_http::test::TestRequest as HttpTestRequest;
 use actix_http::{cookie::Cookie, ws, Extensions, HttpService, Request};
 use actix_router::{Path, ResourceDef, Url};
-use actix_rt::System;
-use actix_server::Server;
+use actix_rt::{time::delay_for, System};
 use actix_service::{
     map_config, IntoService, IntoServiceFactory, Service, ServiceFactory,
 };
@@ -30,7 +29,7 @@ pub use actix_http::test::TestBuffer;
 
 use crate::config::AppConfig;
 use crate::data::Data;
-use crate::dev::{Body, MessageBody, Payload};
+use crate::dev::{Body, MessageBody, Payload, Server};
 use crate::request::HttpRequestPool;
 use crate::rmap::ResourceMap;
 use crate::service::{ServiceRequest, ServiceResponse};
@@ -627,7 +626,7 @@ where
         let ctimeout = cfg.client_timeout;
         let builder = Server::build().workers(1).disable_signals();
 
-        match cfg.stream {
+        let srv = match cfg.stream {
             StreamType::Tcp => match cfg.tp {
                 HttpVer::Http1 => builder.listen("test", tcp, move || {
                     let cfg =
@@ -712,11 +711,11 @@ where
         .unwrap()
         .start();
 
-        tx.send((System::current(), local_addr)).unwrap();
+        tx.send((System::current(), srv, local_addr)).unwrap();
         sys.run()
     });
 
-    let (system, addr) = rx.recv().unwrap();
+    let (system, server, addr) = rx.recv().unwrap();
 
     let client = {
         let connector = {
@@ -752,6 +751,7 @@ where
         addr,
         client,
         system,
+        server,
     }
 }
 
@@ -848,6 +848,7 @@ pub struct TestServer {
     client: awc::Client,
     system: actix_rt::System,
     ssl: bool,
+    server: Server,
 }
 
 impl TestServer {
@@ -936,15 +937,17 @@ impl TestServer {
         self.ws_at("/").await
     }
 
-    /// Stop http server
-    fn stop(&mut self) {
+    /// Gracefully stop http server
+    pub async fn stop(self) {
+        self.server.stop(true).await;
         self.system.stop();
+        delay_for(time::Duration::from_millis(100)).await;
     }
 }
 
 impl Drop for TestServer {
     fn drop(&mut self) {
-        self.stop()
+        self.system.stop()
     }
 }