1 /*
2  * Copyright (C) 2016 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.android.internal.util;
18 
19 import static org.junit.Assert.assertThrows;
20 
21 import android.os.SystemClock;
22 import android.text.format.DateUtils;
23 
24 import junit.framework.TestCase;
25 
26 public class TokenBucketTest extends TestCase {
27 
28     static final int FILL_DELTA_VERY_SHORT  = 1;
29     static final int FILL_DELTA_VERY_LONG   = Integer.MAX_VALUE;
30 
testArgumentValidation()31     public void testArgumentValidation() {
32         assertThrow(() -> new TokenBucket(0, 1, 1));
33         assertThrow(() -> new TokenBucket(1, 0, 1));
34         assertThrow(() -> new TokenBucket(1, 1, 0));
35         assertThrow(() -> new TokenBucket(0, 1));
36         assertThrow(() -> new TokenBucket(1, 0));
37         assertThrow(() -> new TokenBucket(-1, 1, 1));
38         assertThrow(() -> new TokenBucket(1, -1, 1));
39         assertThrow(() -> new TokenBucket(1, 1, -1));
40         assertThrow(() -> new TokenBucket(-1, 1));
41         assertThrow(() -> new TokenBucket(1, -1));
42 
43         new TokenBucket(1000, 100, 0);
44         new TokenBucket(1000, 100, 10);
45         new TokenBucket(5000, 50);
46         new TokenBucket(5000, 1);
47     }
48 
testInitialCapacity()49     public void testInitialCapacity() {
50         drain(new TokenBucket(FILL_DELTA_VERY_LONG, 1), 1);
51         drain(new TokenBucket(FILL_DELTA_VERY_LONG, 10), 10);
52         drain(new TokenBucket(FILL_DELTA_VERY_LONG, 1000), 1000);
53 
54         drain(new TokenBucket(FILL_DELTA_VERY_LONG, 10, 0), 0);
55         drain(new TokenBucket(FILL_DELTA_VERY_LONG, 10, 3), 3);
56         drain(new TokenBucket(FILL_DELTA_VERY_LONG, 10, 10), 10);
57 
58         drain(new TokenBucket(FILL_DELTA_VERY_LONG, 10, 100), 10);
59 
60         drain(new TokenBucket((int) DateUtils.MINUTE_IN_MILLIS, 50), 50);
61         drain(new TokenBucket((int) DateUtils.HOUR_IN_MILLIS, 10), 10);
62         drain(new TokenBucket((int) DateUtils.DAY_IN_MILLIS, 200), 200);
63     }
64 
testReset()65     public void testReset() {
66         TokenBucket tb = new TokenBucket(FILL_DELTA_VERY_LONG, 100, 10);
67         drain(tb, 10);
68 
69         tb.reset(50);
70         drain(tb, 50);
71 
72         tb.reset(50);
73         getOneByOne(tb, 10);
74         assertTrue(tb.has());
75 
76         tb.reset(30);
77         drain(tb, 30);
78     }
79 
testFill()80     public void testFill() throws Exception {
81         int delta = 50;
82         TokenBucket tb = new TokenBucket(delta, 10, 0);
83 
84         assertEmpty(tb);
85 
86         Thread.sleep(3 * delta / 2);
87 
88         assertTrue(tb.has());
89     }
90 
testRefill()91     public void testRefill() throws Exception {
92         TokenBucket tb = new TokenBucket(FILL_DELTA_VERY_SHORT, 10, 10);
93 
94         assertEquals(5, tb.get(5));
95         assertEquals(5, tb.get(5));
96 
97         while (tb.available() < 10) {
98             Thread.sleep(2);
99         }
100 
101         assertEquals(10, tb.get(10));
102 
103         while (tb.available() < 10) {
104             Thread.sleep(2);
105         }
106 
107         assertEquals(10, tb.get(100));
108     }
109 
testAverage()110     public void testAverage() throws Exception {
111         final int delta = 3;
112         final int want = 60;
113 
114         long start = SystemClock.elapsedRealtime();
115         TokenBucket tb = new TokenBucket(delta, 20, 0);
116 
117         for (int i = 0; i < want; i++) {
118             while (!tb.has()) {
119                 Thread.sleep(5 * delta);
120             }
121             tb.get();
122         }
123 
124         assertDuration(want * delta, SystemClock.elapsedRealtime() - start);
125     }
126 
testBurst()127     public void testBurst() throws Exception {
128         final int delta = 2;
129         final int capacity = 20;
130         final int want = 100;
131 
132         long start = SystemClock.elapsedRealtime();
133         TokenBucket tb = new TokenBucket(delta, capacity, 0);
134 
135         int total = 0;
136         while (total < want) {
137             while (!tb.has()) {
138                 Thread.sleep(capacity * delta - 2);
139             }
140             total += tb.get(tb.available());
141         }
142 
143         assertDuration(total * delta, SystemClock.elapsedRealtime() - start);
144     }
145 
getOneByOne(TokenBucket tb, int n)146     static void getOneByOne(TokenBucket tb, int n) {
147         while (n > 0) {
148             assertTrue(tb.has());
149             assertTrue(tb.available() >= n);
150             assertTrue(tb.get());
151             assertTrue(tb.available() >= n - 1);
152             n--;
153         }
154     }
155 
assertEmpty(TokenBucket tb)156     void assertEmpty(TokenBucket tb) {
157         assertFalse(tb.has());
158         assertEquals(0, tb.available());
159         assertFalse(tb.get());
160     }
161 
drain(TokenBucket tb, int n)162     void drain(TokenBucket tb, int n) {
163         getOneByOne(tb, n);
164         assertEmpty(tb);
165     }
166 
assertDuration(long expected, long elapsed)167     void assertDuration(long expected, long elapsed) {
168         String msg = String.format(
169                 "expected elapsed time at least %d ms, but was %d ms", expected, elapsed);
170         elapsed += 1; // one millisecond extra guard
171         assertTrue(msg, elapsed >= expected);
172     }
173 
assertThrow(Fn fn)174     void assertThrow(Fn fn) {
175         assertThrows(Throwable.class, () -> {
176             fn.call();
177         });
178     }
179 
call()180     interface Fn { void call(); }
181 }
182