Skip to content

Commit

Permalink
cleanup and add new test
Browse files Browse the repository at this point in the history
  • Loading branch information
SergioBenitez committed Aug 18, 2024
1 parent 6ce8280 commit 4b2142a
Showing 1 changed file with 27 additions and 30 deletions.
57 changes: 27 additions & 30 deletions core/lib/tests/drop-in-async-context.rs
Original file line number Diff line number Diff line change
@@ -1,45 +1,41 @@
#[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 {
#[launch]
fn launch() -> _ {
super::rocket()
}

#[test]
fn test_launch() {
main();
Expand All @@ -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();
}
}

0 comments on commit 4b2142a

Please sign in to comment.