Fix api::start_stop::test_start_stop test (#136)

Author: mateuszmlc

polybase/tests/api/mod.rs

@@ -103,6 +103,7 @@ impl PortPool {
 #[derive(Debug, Default)]
 struct ServerConfig {
     whitelist: Option<Vec<String>>,
+    keep_port_after_drop: bool,
 }
 
 #[derive(Debug)]
@@ -111,14 +112,17 @@ struct Server {
     // Keep the root dir alive so that polybase can use it
     _root_dir: tempfile::TempDir,
     api_port: u16,
+    keep_port_after_drop: bool,
     client: reqwest::Client,
     base_url: reqwest::Url,
 }
 
 impl Drop for Server {
     fn drop(&mut self) {
         self.process.kill().expect("Failed to stop polybase");
-        API_PORT_POOL.lock().unwrap().release(self.api_port);
+        if !self.keep_port_after_drop {
+            API_PORT_POOL.lock().unwrap().release(self.api_port);
+        }
     }
 }
 
@@ -130,8 +134,10 @@ impl Server {
 
         let mut command = Command::new(find_binary());
 
-        if let Some(whitelist) = config.and_then(|c| c.whitelist) {
-            command.arg("--whitelist").arg(whitelist.join(","));
+        if let Some(ref config) = config {
+            if let Some(ref whitelist) = config.whitelist {
+                command.arg("--whitelist").arg(whitelist.join(","));
+            }
         }
 
         command.arg("--root-dir").arg(root_dir.path());
@@ -155,6 +161,7 @@ impl Server {
             _root_dir: root_dir,
             client: reqwest::Client::new(),
             base_url: format!("http://localhost:{api_port}").parse().unwrap(),
+            keep_port_after_drop: config.map(|c| c.keep_port_after_drop).unwrap_or(false),
             api_port,
         })
     }

polybase/tests/api/start_stop.rs

@@ -1,14 +1,29 @@
+use crate::api::ServerConfig;
+
 #[tokio::test]
 async fn test_start_stop() {
     let api_port;
     {
-        let server = super::Server::setup_and_wait(None).await;
+        let server = super::Server::setup_and_wait(Some(ServerConfig {
+            keep_port_after_drop: true,
+            ..Default::default()
+        }))
+        .await;
         api_port = server.api_port;
     }
 
-    tokio::time::sleep(tokio::time::Duration::from_secs(3)).await;
+    for i in 0..10 {
+        tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
 
-    reqwest::get(format!("http://localhost:{api_port}/v0/health"))
-        .await
-        .expect_err("Server should be stopped");
+        if reqwest::get(format!("http://localhost:{api_port}/v0/health"))
+            .await
+            .is_ok()
+        {
+            if i == 9 {
+                panic!("Server should be stopped");
+            }
+        } else {
+            break;
+        }
+    }
 }

polybase/tests/api/whitelist.rs

@@ -28,7 +28,11 @@ collection Account {
         signature
     });
 
-    let server = Server::setup_and_wait(Some(ServerConfig { whitelist })).await;
+    let server = Server::setup_and_wait(Some(ServerConfig {
+        whitelist,
+        ..Default::default()
+    }))
+    .await;
 
     let res = server
         .create_collection::<Account>("test/Account", schema, Some(&signer))
@@ -62,7 +66,11 @@ collection Account {
         signature
     });
 
-    let server = Server::setup_and_wait(Some(ServerConfig { whitelist })).await;
+    let server = Server::setup_and_wait(Some(ServerConfig {
+        whitelist,
+        ..Default::default()
+    }))
+    .await;
 
     let res = server
         .create_collection::<Account>("test/Account", schema, Some(&signer))
@@ -95,7 +103,11 @@ collection Account {
     let public_key = indexer::PublicKey::from_secp256k1_key(&public_key).unwrap();
     let whitelist = Some(vec![public_key.to_hex().unwrap()]);
 
-    let server = Server::setup_and_wait(Some(ServerConfig { whitelist })).await;
+    let server = Server::setup_and_wait(Some(ServerConfig {
+        whitelist,
+        ..Default::default()
+    }))
+    .await;
 
     let res = server
         .create_collection::<Account>("test/Account", schema, None)