/**
 * 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();
                });
    }

}

@@ -928,128 +928,99 @@
     
-for (int period = 1; period %3C= 5; period++) %7B%0A            for (int capacity = 1; capacity %3C= 5; capacity++) %7B%0A  
+testRateLimitWithGreedyRefill(3, 10);%0A        testRateLimitWithGreedyRefill(3, 3);%0A
         
     
@@ -1007,38 +1007,32 @@
 (3, 3);%0A        
-      
 testRateLimitWit
@@ -1049,50 +1049,14 @@
 ill(
-capacity, period);%0A            %7D%0A        %7D
+4, 2);
 %0A   
@@ -2097,127 +2097,103 @@
     
-for (int period = 1; period %3C= 3; period++) %7B%0A            for (int capacity = 1; capacity %3C= 3; capacity++) %7B%0A 
+testRateLimitWithIntervalRefill(10, 5);%0A        testRateLimitWithIntervalRefill(3, 3);%0A
         
     
@@ -2188,23 +2188,16 @@
         
-       
 testRate
@@ -2224,50 +2224,14 @@
 ill(
-capacity, period);%0A            %7D%0A        %7D
+4, 2);
 %0A   

Fix RateLimitsTest