1 //! This module provides utilities relating to async tasks, typically for usage
2 //! only in test
3
4 use std::{
5 future::{Future, IntoFuture},
6 time::Duration,
7 };
8
9 use bt_common::init_logging;
10 use tokio::{
11 runtime::Builder,
12 select,
13 task::{spawn_local, LocalSet},
14 };
15
16 /// Run the supplied future on a single-threaded runtime
block_on_locally<T>(f: impl Future<Output = T>) -> T17 pub fn block_on_locally<T>(f: impl Future<Output = T>) -> T {
18 init_logging();
19 LocalSet::new().block_on(
20 &Builder::new_current_thread().enable_time().start_paused(true).build().unwrap(),
21 async move {
22 select! {
23 t = f => t,
24 // NOTE: this time should be LARGER than any meaningful delay in the stack
25 _ = tokio::time::sleep(Duration::from_secs(100000)) => {
26 panic!("test appears to be stuck");
27 },
28 }
29 },
30 )
31 }
32
33 /// Check if the supplied future immediately resolves.
34 /// Returns Ok(T) if it resolves, or Err(JoinHandle<T>) if it does not.
35 /// Correctly handles spurious wakeups (unlike Future::poll).
36 ///
37 /// Unlike spawn/spawn_local, try_await guarantees that the future has been
38 /// polled when it returns. In addition, it is safe to drop the returned future,
39 /// since the underlying future will still run (i.e. it will not be cancelled).
40 ///
41 /// Thus, this is useful in tests where we want to force a particular order of
42 /// events, rather than letting spawn_local enqueue a task to the executor at
43 /// *some* point in the future.
44 ///
45 /// MUST only be run in an environment where time is mocked.
try_await<T: 'static>( f: impl IntoFuture<Output = T> + 'static, ) -> Result<T, impl Future<Output = T>>46 pub async fn try_await<T: 'static>(
47 f: impl IntoFuture<Output = T> + 'static,
48 ) -> Result<T, impl Future<Output = T>> {
49 let mut handle = spawn_local(f.into_future());
50
51 select! {
52 t = &mut handle => Ok(t.unwrap()),
53 // NOTE: this time should be SMALLER than any meaningful delay in the stack
54 // since time is frozen in test, we don't need to worry about racing with anything
55 _ = tokio::time::sleep(Duration::from_millis(10)) => {
56 Err(async { handle.await.unwrap() })
57 },
58 }
59 }
60