dendi-wang/thingsboard
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Added missing test for RateLimitServiceTest && refactoring

Browse Source
This commit is contained in:
dshvaika 2025-05-20 16:31:01 +03:00
parent 8af37beb4a
commit a242d3a43b
4 changed files with 71 additions and 92 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
application/src/test/java/org/thingsboard/server/service/limits/RateLimitServiceTest.java
View File

@ -94,6 +94,13 @@ public class RateLimitServiceTest {
testRateLimits(limitedApi, max, tenantId);
}
for (LimitedApi limitedApi : List.of(
LimitedApi.CASSANDRA_READ_QUERIES_MONOLITH,
LimitedApi.CASSANDRA_WRITE_QUERIES_MONOLITH
)) {
testRateLimits(limitedApi, max * 2, tenantId);
}
CustomerId customerId = new CustomerId(UUID.randomUUID());
testRateLimits(LimitedApi.REST_REQUESTS_PER_CUSTOMER, max, customerId);

4
common/data/src/main/java/org/thingsboard/server/common/data/limit/LimitedApiEntry.java
View File

@ -22,10 +22,6 @@ public record LimitedApiEntry(long capacity, long durationSeconds) {
return new LimitedApiEntry(Long.parseLong(parts[0]), Long.parseLong(parts[1]));
}
public double rps() {
return (double) capacity / durationSeconds;
}
@Override
public String toString() {
return capacity + ":" + durationSeconds;

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

@ -1,86 +0,0 @@
/**
* Copyright © 2016-2025 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.awaitility.pollinterval.FixedPollInterval;
import org.junit.jupiter.api.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() {
testRateLimitWithGreedyRefill(3, 10);
testRateLimitWithGreedyRefill(3, 3);
testRateLimitWithGreedyRefill(4, 2);
}
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 = 500;
for (int i = 0; i < capacity; i++) {
await("token refill for rate limit " + rateLimitConfig)
.pollInterval(new FixedPollInterval(10, TimeUnit.MILLISECONDS))
.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() {
testRateLimitWithIntervalRefill(10, 5);
testRateLimitWithIntervalRefill(3, 3);
testRateLimitWithIntervalRefill(4, 2);
}
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 = 500;
await("tokens refill for rate limit " + rateLimitConfig)
.pollInterval(new FixedPollInterval(10, TimeUnit.MILLISECONDS))
.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();
});
}
}

64
common/message/src/test/java/org/thingsboard/server/common/msg/tools/TbRateLimitsTest.java
View File

@ -15,13 +15,75 @@
*/
package org.thingsboard.server.common.msg.tools;
import org.awaitility.pollinterval.FixedPollInterval;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Test;
import java.util.concurrent.TimeUnit;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.awaitility.Awaitility.await;
class TbRateLimitsTest {
public class TbRateLimitsTest {
@Test
public void testRateLimits_greedyRefill() {
testRateLimitWithGreedyRefill(3, 10);
testRateLimitWithGreedyRefill(3, 3);
testRateLimitWithGreedyRefill(4, 2);
}
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 = 500;
for (int i = 0; i < capacity; i++) {
await("token refill for rate limit " + rateLimitConfig)
.pollInterval(new FixedPollInterval(10, TimeUnit.MILLISECONDS))
.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() {
testRateLimitWithIntervalRefill(10, 5);
testRateLimitWithIntervalRefill(3, 3);
testRateLimitWithIntervalRefill(4, 2);
}
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 = 500;
await("tokens refill for rate limit " + rateLimitConfig)
.pollInterval(new FixedPollInterval(10, TimeUnit.MILLISECONDS))
.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();
});
}
@Test
@DisplayName("TbRateLimits should construct with single rate limit")