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