From 922436d38b6d2b57cb5e7de73ee9434e7a84d8fe Mon Sep 17 00:00:00 2001 From: Viacheslav Klimov Date: Tue, 29 Mar 2022 16:30:57 +0300 Subject: [PATCH] Tests for rate limits with different refill strategy --- .../common/msg/tools/RateLimitsTest.java | 87 +++++++++++++++++++ 1 file changed, 87 insertions(+) create mode 100644 common/message/src/test/java/org/thingsboard/server/common/msg/tools/RateLimitsTest.java diff --git a/common/message/src/test/java/org/thingsboard/server/common/msg/tools/RateLimitsTest.java b/common/message/src/test/java/org/thingsboard/server/common/msg/tools/RateLimitsTest.java new file mode 100644 index 0000000000..d4cc408558 --- /dev/null +++ b/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(); + }); + } + +}