diff --git a/Cargo.toml b/Cargo.toml index eae351f..0eb00fe 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,5 +18,5 @@ tracing = "0.1" [dev-dependencies] hyper = { version = "0.14", features = [ "server", "tcp", "http1", "http2" ] } -tokio = { version = "1", features = ["net", "rt-multi-thread", "io-util"] } +tokio = { version = "1", features = ["net", "rt-multi-thread", "io-util", "test-util"] } tracing-subscriber = { version = "0.3", features = ["env-filter"] } diff --git a/README.md b/README.md index 5dc46f3..dfdc988 100644 --- a/README.md +++ b/README.md @@ -35,7 +35,7 @@ and get a sense on how it works, or at least its flow is. > [examples/tokio_tcp.rs](https://github.com/plabayo/tokio-graceful/tree/main/examples/tokio_tcp.rs) > -> ``` +> ```bash > RUST_LOG=trace cargo run --example tokio_tcp > ``` @@ -48,7 +48,7 @@ telnet: `telnet 127.0.0.1 8080`. As you are in control of when to exit you can e > [examples/hyper.rs](https://github.com/plabayo/tokio-graceful/tree/main/examples/hyper.rs) > -> ``` +> ```bash > RUST_LOG=trace cargo run --example hyper > ``` diff --git a/src/lib.rs b/src/lib.rs index 1406e62..aad10d2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,3 +7,39 @@ mod shutdown; pub use shutdown::Shutdown; pub(crate) mod trigger; + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use tokio::sync::oneshot; + + use super::*; + + #[tokio::test] + async fn test_shutdown_nope() { + let (tx, rx) = oneshot::channel::<()>(); + let shutdown = Shutdown::new(async { + rx.await.unwrap(); + }); + tokio::spawn(async move { + tx.send(()).unwrap(); + }); + shutdown.shutdown().await; + } + + #[tokio::test] + async fn test_shutdown_nope_limit() { + let (tx, rx) = oneshot::channel::<()>(); + let shutdown = Shutdown::new(async { + rx.await.unwrap(); + }); + tokio::spawn(async move { + tx.send(()).unwrap(); + }); + shutdown + .shutdown_with_limit(Duration::from_secs(60)) + .await + .unwrap(); + } +}