From bd3715cbd66b2b782a42741782c7ae2a774ce472 Mon Sep 17 00:00:00 2001
From: Sergio Benitez <sb@sergio.bz>
Date: Sat, 17 Aug 2024 18:43:10 -0700
Subject: [PATCH] cleanup and add new test

---
 core/lib/tests/drop-in-async-context.rs | 57 ++++++++++++-------------
 1 file changed, 27 insertions(+), 30 deletions(-)

diff --git a/core/lib/tests/drop-in-async-context.rs b/core/lib/tests/drop-in-async-context.rs
index e1dda6a0..d774f5e3 100644
--- a/core/lib/tests/drop-in-async-context.rs
+++ b/core/lib/tests/drop-in-async-context.rs
@@ -1,38 +1,33 @@
 #[macro_use]
 extern crate rocket;
 
-use rocket::figment::{providers::{Format as _, Toml}, Figment};
-use rocket::{custom, fairing::AdHoc, Build, Orbit, Rocket};
+use rocket::{Build, Config, Rocket};
+use rocket::fairing::AdHoc;
+use rocket::figment::Figment;
 
 struct AsyncDropInAsync;
 
 impl Drop for AsyncDropInAsync {
     fn drop(&mut self) {
-        // Attempt to fetch the current runtime while dropping
-        // Pools in rocket_sync_db_pools (and maybe rocket_db_pools)
-        // do use this capability. They spawn tasks to asyncronously
-        // complete shutdown of the pool, which triggers the same panic.
+        // Ensure that managed state is dropped inside of an async context by
+        // ensuring that we do not panic when fetching the current runtime.
+        //
+        // Crates like rocket_sync_db_pools spawn tasks to asynchronously
+        // complete pool shutdown which must be done in an async context or else
+        // the spawn will panic. We want to ensure that does not happen.
         let _ = rocket::tokio::runtime::Handle::current();
     }
 }
 
 fn rocket() -> Rocket<Build> {
-    let mut config = rocket::Config::default();
-    #[cfg(feature = "secrets")]
-    { config.secret_key = rocket::config::SecretKey::generate().unwrap(); }
-    let figment = Figment::from(config).merge(Toml::string(r#"
-[default]
-address = "tcp:127.0.0.1:0"
-port = 0
-"#).nested());
-    custom(figment).manage(AsyncDropInAsync).attach(AdHoc::on_liftoff(
-        "Shutdown immediately",
-        |rocket: &Rocket<Orbit>| {
-            Box::pin(async {
-                rocket.shutdown().notify();
-            })
-        },
-    ))
+    let figment = Figment::from(Config::debug_default())
+        .merge(("address", "tcp:127.0.0.1:0"));
+
+    rocket::custom(figment)
+        .manage(AsyncDropInAsync)
+        .attach(AdHoc::on_liftoff("Shutdown", |rocket| Box::pin(async {
+            rocket.shutdown().notify();
+        })))
 }
 
 mod launch {
@@ -40,6 +35,7 @@ mod launch {
     fn launch() -> _ {
         super::rocket()
     }
+
     #[test]
     fn test_launch() {
         main();
@@ -49,22 +45,23 @@ mod launch {
 mod main {
     #[rocket::main]
     async fn main() {
-        super::rocket()
-            .launch()
-            .await
-            .unwrap();
+        super::rocket().launch().await.unwrap();
     }
+
     #[test]
     fn test_main() {
         main();
     }
+
     #[test]
     fn test_execute() {
         rocket::execute(async {
-            super::rocket()
-                .launch()
-                .await
-                .unwrap();
+            super::rocket().launch().await.unwrap();
         });
     }
+
+    #[test]
+    fn test_execute_directly() {
+        rocket::execute(super::rocket().launch()).unwrap();
+    }
 }