From c10ba684e0e22250b55bef0d14f0898ac900b28d Mon Sep 17 00:00:00 2001
From: Conrad Ludgate <conradludgate@gmail.com>
Date: Thu, 27 Jul 2023 09:34:13 +0100
Subject: some simple server tests (#1096)

* some simple server tests

* fmt

* logging in server test

* log server errors

* fix postgres uri

* postgres ports

* localhost again?

* Rebase fixes

---------

Co-authored-by: Ellie Huxtable <ellie@elliehuxtable.com>
---
 atuin-server/src/lib.rs | 35 ++++++++++++++++++++++++++++-------
 1 file changed, 28 insertions(+), 7 deletions(-)

(limited to 'atuin-server/src/lib.rs')

diff --git a/atuin-server/src/lib.rs b/atuin-server/src/lib.rs
index fc4d9e0c..810b7db9 100644
--- a/atuin-server/src/lib.rs
+++ b/atuin-server/src/lib.rs
@@ -1,6 +1,9 @@
 #![forbid(unsafe_code)]
 
-use std::net::{IpAddr, SocketAddr};
+use std::{
+    future::Future,
+    net::{IpAddr, SocketAddr, TcpListener},
+};
 
 use atuin_server_database::Database;
 use axum::Server;
@@ -16,10 +19,15 @@ use tokio::signal;
 
 #[cfg(target_family = "unix")]
 async fn shutdown_signal() {
-    signal::unix::signal(signal::unix::SignalKind::terminate())
-        .expect("failed to register signal handler")
-        .recv()
-        .await;
+    let mut term = signal::unix::signal(signal::unix::SignalKind::terminate())
+        .expect("failed to register signal handler");
+    let mut interrupt = signal::unix::signal(signal::unix::SignalKind::interrupt())
+        .expect("failed to register signal handler");
+
+    tokio::select! {
+        _ = term.recv() => {},
+        _ = interrupt.recv() => {},
+    };
     eprintln!("Shutting down gracefully...");
 }
 
@@ -38,16 +46,29 @@ pub async fn launch<Db: Database>(
     port: u16,
 ) -> Result<()> {
     let host = host.parse::<IpAddr>()?;
+    launch_with_listener::<Db>(
+        settings,
+        TcpListener::bind(SocketAddr::new(host, port)).context("could not connect to socket")?,
+        shutdown_signal(),
+    )
+    .await
+}
 
+pub async fn launch_with_listener<Db: Database>(
+    settings: Settings<Db::Settings>,
+    listener: TcpListener,
+    shutdown: impl Future<Output = ()>,
+) -> Result<()> {
     let db = Db::new(&settings.db_settings)
         .await
         .wrap_err_with(|| format!("failed to connect to db: {:?}", settings.db_settings))?;
 
     let r = router::router(db, settings);
 
-    Server::bind(&SocketAddr::new(host, port))
+    Server::from_tcp(listener)
+        .context("could not launch server")?
         .serve(r.into_make_service())
-        .with_graceful_shutdown(shutdown_signal())
+        .with_graceful_shutdown(shutdown)
         .await?;
 
     Ok(())
-- 
cgit v1.3.1