From 2deab8429c52193cb0edc65cee9825c62c2c9750 Mon Sep 17 00:00:00 2001
From: Salvador Girones Gil <salvadorgirones@gmail.com>
Date: Mon, 21 Aug 2023 14:58:53 +0200
Subject: [PATCH] [cli/server] Add tests for signal handling (#2481)

---
 tests/cli_integration.rs | 90 ++++++++++++++++++++++++++++++++++++++++
 tests/common/mod.rs      | 11 +++++
 2 files changed, 101 insertions(+)

diff --git a/tests/cli_integration.rs b/tests/cli_integration.rs
index 86460428..b20e814f 100644
--- a/tests/cli_integration.rs
+++ b/tests/cli_integration.rs
@@ -414,3 +414,93 @@ fn validate_failed_due_to_invalid_surql_files_syntax() {
 
 	assert!(common::run_in_dir("validate", &temp_dir).output().is_err());
 }
+
+#[test(tokio::test)]
+async fn test_server_graceful_shutdown() {
+	let (_, mut server) = common::start_server(StartServerArguments {
+		auth: false,
+		tls: false,
+		wait_is_ready: true,
+		tick_interval: ONE_SEC,
+	})
+	.await
+	.unwrap();
+
+	info!("* Send SIGINT signal");
+	server.send_signal(nix::sys::signal::Signal::SIGINT).expect("Failed to send SIGINT to server");
+
+	info!("* Waiting for server to exit gracefully ...");
+	tokio::select! {
+		_ = async {
+			loop {
+				if let Ok(Some(exit)) = server.status() {
+					assert!(exit.success(), "Server shutted down successfully");
+					break;
+				}
+				tokio::time::sleep(time::Duration::from_secs(1)).await;
+			}
+		} => {},
+		// Timeout after 5 seconds
+		_ = tokio::time::sleep(time::Duration::from_secs(5)) => {
+			panic!("Server didn't exit after receiving SIGINT");
+		}
+	}
+}
+
+#[test(tokio::test)]
+async fn test_server_second_signal_handling() {
+	let (addr, mut server) = common::start_server(StartServerArguments {
+		auth: false,
+		tls: false,
+		wait_is_ready: true,
+		tick_interval: ONE_SEC,
+	})
+	.await
+	.unwrap();
+
+	// Create a long-lived WS connection so the server don't shutdown gracefully
+	let mut socket = common::connect_ws(&addr).await.expect("Failed to connect to server");
+	let json = serde_json::json!({
+		"id": "1",
+		"method": "query",
+		"params": ["SLEEP 30s;"],
+	});
+	common::ws_send_msg(&mut socket, serde_json::to_string(&json).unwrap())
+		.await
+		.expect("Failed to send WS message");
+
+	info!("* Send first SIGINT signal");
+	server.send_signal(nix::sys::signal::Signal::SIGINT).expect("Failed to send SIGINT to server");
+
+	tokio::select! {
+		_ = async {
+			loop {
+				if let Ok(Some(exit)) = server.status() {
+					panic!("Server unexpectedly exited after receiving first SIGINT: {:?}", exit);
+				}
+				tokio::time::sleep(time::Duration::from_secs(1)).await;
+			}
+		} => {},
+		// Timeout after 5 seconds
+		_ = tokio::time::sleep(time::Duration::from_secs(5)) => ()
+	}
+
+	info!("* Send second SIGINT signal");
+	server.send_signal(nix::sys::signal::Signal::SIGINT).expect("Failed to send SIGINT to server");
+
+	tokio::select! {
+		_ = async {
+			loop {
+				if let Ok(Some(exit)) = server.status() {
+					assert!(exit.success(), "Server shutted down successfully");
+					break;
+				}
+				tokio::time::sleep(time::Duration::from_secs(1)).await;
+			}
+		} => {},
+		// Timeout after 5 seconds
+		_ = tokio::time::sleep(time::Duration::from_secs(5)) => {
+			panic!("Server didn't exit after receiving two SIGINT signals");
+		}
+	}
+}
diff --git a/tests/common/mod.rs b/tests/common/mod.rs
index c9826b4c..81715573 100644
--- a/tests/common/mod.rs
+++ b/tests/common/mod.rs
@@ -43,6 +43,17 @@ impl Child {
 		self
 	}
 
+	pub fn send_signal(&self, signal: nix::sys::signal::Signal) -> nix::Result<()> {
+		nix::sys::signal::kill(
+			nix::unistd::Pid::from_raw(self.inner.as_ref().unwrap().id() as i32),
+			signal,
+		)
+	}
+
+	pub fn status(&mut self) -> std::io::Result<Option<std::process::ExitStatus>> {
+		self.inner.as_mut().unwrap().try_wait()
+	}
+
 	/// Read the child's stdout concatenated with its stderr. Returns Ok if the child
 	/// returns successfully, Err otherwise.
 	pub fn output(mut self) -> Result<String, String> {