Tests for rate limits with different refill strategy

common/message/src/test/java/org/thingsboard/server/common/msg/tools/RateLimitsTest.java

@@ -0,0 +1,87 @@
+/**
+ * Copyright © 2016-2022 The Thingsboard Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.thingsboard.server.common.msg.tools;
+
+import org.junit.Test;
+
+import java.util.concurrent.TimeUnit;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.awaitility.Awaitility.await;
+
+public class RateLimitsTest {
+
+    @Test
+    public void testRateLimits_greedyRefill() {
+        for (int period = 1; period <= 5; period++) {
+            for (int capacity = 1; capacity <= 5; capacity++) {
+                testRateLimitWithGreedyRefill(capacity, period);
+            }
+        }
+    }
+
+    private void testRateLimitWithGreedyRefill(int capacity, int period) {
+        String rateLimitConfig = capacity + ":" + period;
+        TbRateLimits rateLimits = new TbRateLimits(rateLimitConfig);
+
+        rateLimits.tryConsume(capacity);
+        assertThat(rateLimits.tryConsume()).as("new token is available").isFalse();
+
+        int expectedRefillTime = (int) (((double) period / capacity) * 1000);
+        int gap = 100;
+
+        for (int i = 0; i < capacity; i++) {
+            await("token refill for rate limit " + rateLimitConfig)
+                    .atLeast(expectedRefillTime - gap, TimeUnit.MILLISECONDS)
+                    .atMost(expectedRefillTime + gap, TimeUnit.MILLISECONDS)
+                    .untilAsserted(() -> {
+                        assertThat(rateLimits.tryConsume()).as("token is available").isTrue();
+                    });
+            assertThat(rateLimits.tryConsume()).as("new token is available").isFalse();
+        }
+    }
+
+    @Test
+    public void testRateLimits_intervalRefill() {
+        for (int period = 1; period <= 3; period++) {
+            for (int capacity = 1; capacity <= 3; capacity++) {
+                testRateLimitWithIntervalRefill(capacity, period);
+            }
+        }
+    }
+
+    private void testRateLimitWithIntervalRefill(int capacity, int period) {
+        String rateLimitConfig = capacity + ":" + period;
+        TbRateLimits rateLimits = new TbRateLimits(rateLimitConfig, true);
+
+        rateLimits.tryConsume(capacity);
+        assertThat(rateLimits.tryConsume()).as("new token is available").isFalse();
+
+        int expectedRefillTime = period * 1000;
+        int gap = 100;
+
+        await("tokens refill for rate limit " + rateLimitConfig)
+                .atLeast(expectedRefillTime - gap, TimeUnit.MILLISECONDS)
+                .atMost(expectedRefillTime + gap, TimeUnit.MILLISECONDS)
+                .untilAsserted(() -> {
+                    for (int i = 0; i < capacity; i++) {
+                        assertThat(rateLimits.tryConsume()).as("token is available").isTrue();
+                    }
+                    assertThat(rateLimits.tryConsume()).as("new token is available").isFalse();
+                });
+    }
+
+}