RateLimiterTest.java

import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Test;

import java.time.Duration;
import java.time.Instant;
import java.util.UUID;

import static org.assertj.core.api.Assertions.assertThat;

class RateLimiterTest {

    @Test
    @DisplayName("Allows up to window limit")
    void allowsUpToLimit() {
        var clock = new TimeSource(Instant.EPOCH);
        var limiter = new RateLimiter<String>(3, Duration.ofNanos(10_000L), clock);

        assertThat(limiter.allow("A")).isTrue();
        clock.advance(Duration.ofNanos(1L));
        assertThat(limiter.allow("A")).isTrue();
        assertThat(limiter.allow("A")).isTrue();
        assertThat(limiter.allow("A")).isFalse();
    }
    
    @Test
    @DisplayName("Denies requests just before boundary")
    void denyCloseToBoundary() {
        var clock = new TimeSource(Instant.EPOCH);
        var limiter = new RateLimiter<String>(2, Duration.ofNanos(10_000L), clock);

        assertThat(limiter.allow("A")).isTrue();
        assertThat(limiter.allow("A")).isTrue();
        assertThat(limiter.allow("A")).isFalse();

        // Just before boundary: still same window
        clock.advance(Duration.ofNanos(9_999L));
        assertThat(limiter.allow("A")).isFalse();
    }

    @Test
    @DisplayName("Allows first request at exact boundary")
    void allowsNewBoundary() {
        var clock = new TimeSource(Instant.EPOCH);
        var limiter = new RateLimiter<String>(2, Duration.ofNanos(10_000L), clock);

        assertThat(limiter.allow("A")).isTrue();
        assertThat(limiter.allow("A")).isTrue();
        assertThat(limiter.allow("A")).isFalse();

        // At exact boundary: new window
        clock.advance(Duration.ofNanos(10_000L));
        assertThat(limiter.allow("A")).isTrue();
    }

    @Test
    @DisplayName("Resets at boundary, then counts within window")
    void continuesCountingWithinWindowAfterBoundaryReset() {
        var clock = new TimeSource(Instant.EPOCH);
        var limiter = new RateLimiter<String>(2, Duration.ofNanos(5_000L), clock);

        assertThat(limiter.allow("key")).isTrue();
        clock.advance(Duration.ofNanos(1L));
        assertThat(limiter.allow("key")).isTrue();
        assertThat(limiter.allow("key")).isFalse();

        // Jump to next window
        clock.advance(Duration.ofNanos(5_000L));
        assertThat(limiter.allow("key")).isTrue();
        assertThat(limiter.allow("key")).isTrue();
        assertThat(limiter.allow("key")).isFalse();
    }

    @Test
    @DisplayName("Independent counters/windows per key")
    void separateKeysHaveIndependentCountersAndWindows() {
        var clock = new TimeSource(Instant.EPOCH.plusNanos(42L));
        var limiter = new RateLimiter<String>(1, Duration.ofNanos(100L), clock);

        assertThat(limiter.allow("A")).isTrue();
        // Advance a tiny amount to model time moving between requests
        clock.advance(Duration.ofNanos(1L));
        assertThat(limiter.allow("A")).isFalse();
        assertThat(limiter.allow("B")).isTrue(); // independent key
        assertThat(limiter.allow("B")).isFalse();

        clock.advance(Duration.ofNanos(100L)); // new window for both at boundary
        assertThat(limiter.allow("A")).isTrue();
        assertThat(limiter.allow("B")).isTrue();
    }

    @Test
    @DisplayName("Long gaps reset window")
    void longGapsResetWindow() {
        var clock = new TimeSource(Instant.EPOCH.plusNanos(1_000L));
        var limiter = new RateLimiter<String>(2, Duration.ofNanos(50L), clock);

        assertThat(limiter.allow("X")).isTrue();
        clock.advance(Duration.ofNanos(1L));
        assertThat(limiter.allow("X")).isTrue();
        assertThat(limiter.allow("X")).isFalse();

        // Advance several windows worth
        clock.advance(Duration.ofNanos(1_000L));
        assertThat(limiter.allow("X")).isTrue();
        assertThat(limiter.allow("X")).isTrue();
        assertThat(limiter.allow("X")).isFalse();
    }

    @Test
    @DisplayName("Exact boundary starts a new window each time")
    void exactBoundaryIsNewWindowEveryTime() {
        var clock = new TimeSource(Instant.EPOCH);
        var limiter = new RateLimiter<String>(1, Duration.ofNanos(10L), clock);

        assertThat(limiter.allow("k")).isTrue();
        assertThat(limiter.allow("k")).isFalse();

        // Move exactly to boundary repeatedly; each time should allow once
        for (int i = 0; i < 5; i++) {
            clock.advance(Duration.ofNanos(10L));
            assertThat(limiter.allow("k")).isTrue();
            assertThat(limiter.allow("k")).isFalse();
        }
    }

    @Test
    @DisplayName("Supports UUID keys")
    void supportsUuidKeys() {
        var clock = new TimeSource(Instant.EPOCH);
        var limiter = new RateLimiter<UUID>(1, Duration.ofSeconds(1L), clock);

        var a = UUID.fromString("00000000-0000-0000-0000-000000000001");
        var b = UUID.fromString("00000000-0000-0000-0000-000000000002");

        assertThat(limiter.allow(a)).isTrue();
        assertThat(limiter.allow(a)).isFalse();
        // Advance slightly so the second client's window anchors later
        clock.advance(Duration.ofMillis(1L));
        assertThat(limiter.allow(b)).isTrue();
        assertThat(limiter.allow(b)).isFalse();

        // Advance exactly one second to hit the window boundary
        clock.advance(Duration.ofSeconds(1L));
        assertThat(limiter.allow(a)).isTrue();
        assertThat(limiter.allow(b)).isTrue();
    }

    @Test
    @DisplayName("Supports Integer keys")
    void supportsIntegerKeys() {
        var clock = new TimeSource(Instant.EPOCH);
        var limiter = new RateLimiter<Integer>(1, Duration.ofNanos(100L), clock);

        assertThat(limiter.allow(42)).isTrue();
        assertThat(limiter.allow(42)).isFalse();
        clock.advance(Duration.ofNanos(1L));
        assertThat(limiter.allow(84)).isTrue(); // independent key of different type
        assertThat(limiter.allow(84)).isFalse();

        clock.advance(Duration.ofNanos(100L)); // boundary resets both
        assertThat(limiter.allow(42)).isTrue();
        assertThat(limiter.allow(84)).isTrue();
    }
    
    @Test
    @DisplayName("Supports Long keys")
    void supportsLongKeys() {
        var clock = new TimeSource(Instant.EPOCH);
        var limiter = new RateLimiter<Long>(2, Duration.ofNanos(50L), clock);

        assertThat(limiter.allow(1L)).isTrue();
        assertThat(limiter.allow(1L)).isTrue();
        assertThat(limiter.allow(1L)).isFalse();

        clock.advance(Duration.ofNanos(1L));
        assertThat(limiter.allow(2L)).isTrue(); // independent key
        assertThat(limiter.allow(2L)).isTrue();
        assertThat(limiter.allow(2L)).isFalse();

        clock.advance(Duration.ofNanos(50L));
        assertThat(limiter.allow(1L)).isTrue();
        assertThat(limiter.allow(2L)).isTrue();
    }
}