Add cancel_on_shutdown test

tests/cancel_on_shutdown.rs

@@ -0,0 +1,77 @@
+use tokio::time::{sleep, Duration};
+use tokio_graceful_shutdown::{
+    errors::CancelledByShutdown, FutureExt, SubsystemBuilder, SubsystemHandle, Toplevel,
+};
+use tracing_test::traced_test;
+
+use std::error::Error;
+
+/// Wrapper function to simplify lambdas
+type BoxedError = Box<dyn Error + Sync + Send>;
+type BoxedResult = Result<(), BoxedError>;
+
+#[tokio::test]
+#[traced_test]
+async fn cancel_on_shutdown_propagates_result() {
+    let subsystem1 = |subsys: SubsystemHandle| async move {
+        let compute_value = async {
+            sleep(Duration::from_millis(10)).await;
+            42
+        };
+
+        let value = compute_value.cancel_on_shutdown(&subsys).await;
+
+        assert_eq!(value.ok(), Some(42));
+
+        BoxedResult::Ok(())
+    };
+
+    let subsystem2 = |subsys: SubsystemHandle| async move {
+        async fn compute_value() -> i32 {
+            sleep(Duration::from_millis(10)).await;
+            42
+        }
+
+        let value = compute_value().cancel_on_shutdown(&subsys).await;
+
+        assert_eq!(value.ok(), Some(42));
+
+        BoxedResult::Ok(())
+    };
+
+    let result = Toplevel::<BoxedError>::new(move |s| async move {
+        s.start(SubsystemBuilder::new("subsys1", subsystem1));
+        s.start(SubsystemBuilder::new("subsys2", subsystem2));
+    })
+    .handle_shutdown_requests(Duration::from_millis(200))
+    .await;
+
+    assert!(result.is_ok());
+}
+
+#[tokio::test]
+#[traced_test]
+async fn cancel_on_shutdown_cancels_on_shutdown() {
+    let subsystem = |subsys: SubsystemHandle| async move {
+        async fn compute_value(subsys: &SubsystemHandle) -> i32 {
+            sleep(Duration::from_millis(100)).await;
+            subsys.request_shutdown();
+            sleep(Duration::from_millis(100)).await;
+            42
+        }
+
+        let value = compute_value(&subsys).cancel_on_shutdown(&subsys).await;
+
+        assert!(matches!(value, Err(CancelledByShutdown)));
+
+        BoxedResult::Ok(())
+    };
+
+    let result = Toplevel::<BoxedError>::new(move |s| async move {
+        s.start(SubsystemBuilder::new("subsys", subsystem));
+    })
+    .handle_shutdown_requests(Duration::from_millis(200))
+    .await;
+
+    assert!(result.is_ok());
+}