1 use super::support;
2 use std::time::Duration;
3 use tokio::time;
4 use tokio_test::{assert_pending, assert_ready, assert_ready_ok};
5 use tower::limit::rate::RateLimitLayer;
6 use tower_test::{assert_request_eq, mock};
7
8 #[tokio::test(flavor = "current_thread")]
reaching_capacity()9 async fn reaching_capacity() {
10 let _t = support::trace_init();
11 time::pause();
12
13 let rate_limit = RateLimitLayer::new(1, Duration::from_millis(100));
14 let (mut service, mut handle) = mock::spawn_layer(rate_limit);
15
16 assert_ready_ok!(service.poll_ready());
17
18 let response = service.call("hello");
19
20 assert_request_eq!(handle, "hello").send_response("world");
21
22 assert_eq!(response.await.unwrap(), "world");
23 assert_pending!(service.poll_ready());
24
25 assert_pending!(handle.poll_request());
26
27 time::advance(Duration::from_millis(101)).await;
28
29 assert_ready_ok!(service.poll_ready());
30
31 let response = service.call("two");
32
33 assert_request_eq!(handle, "two").send_response("done");
34
35 assert_eq!(response.await.unwrap(), "done");
36 }
37
38 #[tokio::test(flavor = "current_thread")]
remaining_gets_reset()39 async fn remaining_gets_reset() {
40 // This test checks for the case where the `until` state gets reset
41 // but the `rem` does not. This was a bug found `cd7dd12315706fc0860a35646b1eb7b60c50a5c1`.
42 //
43 // The main premise here is that we can make one request which should initialize the state
44 // as ready. Then we can advance the clock to put us beyond the current period. When we make
45 // subsequent requests the `rem` for the next window is continued from the previous when
46 // it should be totally reset.
47 let _t = support::trace_init();
48 time::pause();
49
50 let rate_limit = RateLimitLayer::new(3, Duration::from_millis(100));
51 let (mut service, mut handle) = mock::spawn_layer(rate_limit);
52
53 assert_ready_ok!(service.poll_ready());
54 let response = service.call("hello");
55 assert_request_eq!(handle, "hello").send_response("world");
56 assert_eq!(response.await.unwrap(), "world");
57
58 time::advance(Duration::from_millis(100)).await;
59
60 assert_ready_ok!(service.poll_ready());
61 let response = service.call("hello");
62 assert_request_eq!(handle, "hello").send_response("world");
63 assert_eq!(response.await.unwrap(), "world");
64
65 assert_ready_ok!(service.poll_ready());
66 let response = service.call("hello");
67 assert_request_eq!(handle, "hello").send_response("world");
68 assert_eq!(response.await.unwrap(), "world");
69
70 assert_ready_ok!(service.poll_ready());
71 }
72