From be20e9dcc0f88993be87f5f26b76ae69ef4d8436 Mon Sep 17 00:00:00 2001
From: David Pedersen <david.pdrsn@gmail.com>
Date: Wed, 28 Apr 2021 11:44:06 +0200
Subject: [PATCH] Refactor tests

---
 tests/integration_tests/tests/timeout.rs | 95 ++++++++----------------
 1 file changed, 29 insertions(+), 66 deletions(-)

diff --git a/tests/integration_tests/tests/timeout.rs b/tests/integration_tests/tests/timeout.rs
index 7a4b94e6b..450a67d21 100644
--- a/tests/integration_tests/tests/timeout.rs
+++ b/tests/integration_tests/tests/timeout.rs
@@ -1,33 +1,11 @@
 use integration_tests::pb::{test_client, test_server, Input, Output};
-use std::time::Duration;
+use std::{net::SocketAddr, time::Duration};
 use tokio::net::TcpListener;
 use tonic::{transport::Server, Code, Request, Response, Status};
 
 #[tokio::test]
 async fn cancelation_on_timeout() {
-    struct Svc;
-
-    #[tonic::async_trait]
-    impl test_server::Test for Svc {
-        async fn unary_call(&self, _req: Request<Input>) -> Result<Response<Output>, Status> {
-            // Wait for a time longer than the timeout
-            tokio::time::sleep(Duration::from_millis(1_000)).await;
-            Ok(Response::new(Output {}))
-        }
-    }
-
-    let svc = test_server::TestServer::new(Svc);
-
-    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
-    let addr = listener.local_addr().unwrap();
-
-    tokio::spawn(async move {
-        Server::builder()
-            .add_service(svc)
-            .serve_with_incoming(tokio_stream::wrappers::TcpListenerStream::new(listener))
-            .await
-            .unwrap();
-    });
+    let addr = run_service_in_background(Duration::from_secs(1), Duration::from_secs(100)).await;
 
     let mut client = test_client::TestClient::connect(format!("http://{}", addr))
         .await
@@ -35,6 +13,7 @@ async fn cancelation_on_timeout() {
 
     let mut req = Request::new(Input {});
     req.metadata_mut()
+        // 500 ms
         .insert("grpc-timeout", "500m".parse().unwrap());
 
     let res = client.unary_call(req).await;
@@ -46,30 +25,7 @@ async fn cancelation_on_timeout() {
 
 #[tokio::test]
 async fn picks_server_timeout_if_thats_sorter() {
-    struct Svc;
-
-    #[tonic::async_trait]
-    impl test_server::Test for Svc {
-        async fn unary_call(&self, _req: Request<Input>) -> Result<Response<Output>, Status> {
-            // Wait for a time longer than the timeout
-            tokio::time::sleep(Duration::from_secs(1)).await;
-            Ok(Response::new(Output {}))
-        }
-    }
-
-    let svc = test_server::TestServer::new(Svc);
-
-    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
-    let addr = listener.local_addr().unwrap();
-
-    tokio::spawn(async move {
-        Server::builder()
-            .timeout(Duration::from_millis(100))
-            .add_service(svc)
-            .serve_with_incoming(tokio_stream::wrappers::TcpListenerStream::new(listener))
-            .await
-            .unwrap();
-    });
+    let addr = run_service_in_background(Duration::from_secs(1), Duration::from_millis(100)).await;
 
     let mut client = test_client::TestClient::connect(format!("http://{}", addr))
         .await
@@ -88,42 +44,49 @@ async fn picks_server_timeout_if_thats_sorter() {
 
 #[tokio::test]
 async fn picks_client_timeout_if_thats_sorter() {
-    struct Svc;
+    let addr = run_service_in_background(Duration::from_secs(1), Duration::from_secs(100)).await;
+
+    let mut client = test_client::TestClient::connect(format!("http://{}", addr))
+        .await
+        .unwrap();
+
+    let mut req = Request::new(Input {});
+    req.metadata_mut()
+        // 100 ms
+        .insert("grpc-timeout", "100m".parse().unwrap());
+
+    let res = client.unary_call(req).await;
+    let err = res.unwrap_err();
+    assert!(err.message().contains("Timeout expired"));
+    assert_eq!(err.code(), Code::Cancelled);
+}
+
+async fn run_service_in_background(latency: Duration, server_timeout: Duration) -> SocketAddr {
+    struct Svc {
+        latency: Duration,
+    }
 
     #[tonic::async_trait]
     impl test_server::Test for Svc {
         async fn unary_call(&self, _req: Request<Input>) -> Result<Response<Output>, Status> {
-            // Wait for a time longer than the timeout
-            tokio::time::sleep(Duration::from_secs(1)).await;
+            tokio::time::sleep(self.latency).await;
             Ok(Response::new(Output {}))
         }
     }
 
-    let svc = test_server::TestServer::new(Svc);
+    let svc = test_server::TestServer::new(Svc { latency });
 
     let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
     let addr = listener.local_addr().unwrap();
 
     tokio::spawn(async move {
         Server::builder()
-            .timeout(Duration::from_secs(9001))
+            .timeout(server_timeout)
             .add_service(svc)
             .serve_with_incoming(tokio_stream::wrappers::TcpListenerStream::new(listener))
             .await
             .unwrap();
     });
 
-    let mut client = test_client::TestClient::connect(format!("http://{}", addr))
-        .await
-        .unwrap();
-
-    let mut req = Request::new(Input {});
-    req.metadata_mut()
-        // 100 ms
-        .insert("grpc-timeout", "100m".parse().unwrap());
-
-    let res = client.unary_call(req).await;
-    let err = res.unwrap_err();
-    assert!(err.message().contains("Timeout expired"));
-    assert_eq!(err.code(), Code::Cancelled);
+    addr
 }