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