1  #![cfg(feature = "load-shed")]
2  #[path = "../support.rs"]
3  mod support;
4  
5  use tokio_test::{assert_ready_err, assert_ready_ok, task};
6  use tower::load_shed::LoadShedLayer;
7  use tower_test::{assert_request_eq, mock};
8  
9  #[tokio::test(flavor = "current_thread")]
when_ready()10  async fn when_ready() {
11      let _t = support::trace_init();
12  
13      let layer = LoadShedLayer::new();
14      let (mut service, mut handle) = mock::spawn_layer(layer);
15  
16      assert_ready_ok!(service.poll_ready(), "overload always reports ready");
17  
18      let mut response = task::spawn(service.call("hello"));
19  
20      assert_request_eq!(handle, "hello").send_response("world");
21      assert_eq!(assert_ready_ok!(response.poll()), "world");
22  }
23  
24  #[tokio::test(flavor = "current_thread")]
when_not_ready()25  async fn when_not_ready() {
26      let _t = support::trace_init();
27  
28      let layer = LoadShedLayer::new();
29      let (mut service, mut handle) = mock::spawn_layer::<_, (), _>(layer);
30  
31      handle.allow(0);
32  
33      assert_ready_ok!(service.poll_ready(), "overload always reports ready");
34  
35      let mut fut = task::spawn(service.call("hello"));
36  
37      let err = assert_ready_err!(fut.poll());
38      assert!(err.is::<tower::load_shed::error::Overloaded>());
39  }
40