spring-boot-demo/rate-limit/src/test/java/com/ankurm/ratelimit/PerUserLimitTest.java

package com.ankurm.ratelimit;

import java.time.Duration;

import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.MethodOrderer;
import org.junit.jupiter.api.Order;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.TestMethodOrder;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.test.web.servlet.MockMvc;
import org.springframework.test.web.servlet.MvcResult;

import static org.assertj.core.api.Assertions.assertThat;
import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.get;

/** Per-user limit: 10 requests per minute, greedy refill (one token every 6 s). Writes output/02-per-user.txt. */
@TestMethodOrder(MethodOrderer.OrderAnnotation.class)
class PerUserLimitTest extends RedisBackedTest {

    static Transcript t;

    @Autowired
    MockMvc mvc;

    @BeforeAll
    static void open() {
        t = Transcript.open("02-per-user.txt");
    }

    @AfterAll
    static void close() {
        t.close();
    }

    private MvcResult call(String user) throws Exception {
        MvcResult r = mvc.perform(get("/api/hello").header("X-User-Id", user)).andReturn();
        var res = r.getResponse();
        t.line("GET /api/hello  X-User-Id: %-5s -> %d  limit=%s remaining=%s reset=%s retry-after=%s scope=%s",
                user, res.getStatus(), res.getHeader("X-RateLimit-Limit"), res.getHeader("X-RateLimit-Remaining"),
                res.getHeader("X-RateLimit-Reset"), res.getHeader("Retry-After"), res.getHeader("X-RateLimit-Scope"));
        return r;
    }

    @Test
    @Order(1)
    void eleventhRequestInAMinuteGets429() throws Exception {
        t.line("== alice sends 12 requests with no pause");
        for (int i = 1; i <= 10; i++) {
            var res = call("alice").getResponse();
            assertThat(res.getStatus()).isEqualTo(200);
            assertThat(res.getHeader("X-RateLimit-Remaining")).isEqualTo(String.valueOf(10 - i));
        }
        MvcResult rejected = call("alice");
        assertThat(rejected.getResponse().getStatus()).isEqualTo(429);
        assertThat(rejected.getResponse().getHeader("Retry-After")).isEqualTo("6");
        assertThat(rejected.getResponse().getContentType()).isEqualTo("application/problem+json");
        call("alice");
        t.line("-- body of the 429:");
        t.line(rejected.getResponse().getContentAsString());
    }

    @Test
    @Order(2)
    void aDifferentUserIsUnaffected() throws Exception {
        t.line("== alice is out of tokens; bob is a different bucket");
        for (int i = 0; i < 10; i++) {
            mvc.perform(get("/api/hello").header("X-User-Id", "alice"));
        }
        assertThat(call("alice").getResponse().getStatus()).isEqualTo(429);
        MvcResult bob = call("bob");
        assertThat(bob.getResponse().getStatus()).isEqualTo(200);
        assertThat(bob.getResponse().getHeader("X-RateLimit-Remaining")).isEqualTo("9");
    }

    @Test
    @Order(3)
    void honouringRetryAfterWorks() throws Exception {
        t.line("== alice waits exactly Retry-After (6 s on the test clock) and tries again");
        for (int i = 0; i < 10; i++) {
            mvc.perform(get("/api/hello").header("X-User-Id", "alice"));
        }
        MvcResult rejected = call("alice");
        long wait = Long.parseLong(rejected.getResponse().getHeader("Retry-After"));
        clock.advance(Duration.ofSeconds(wait - 1));
        t.line("-- after %d s: still too early", wait - 1);
        assertThat(call("alice").getResponse().getStatus()).isEqualTo(429);
        clock.advance(Duration.ofSeconds(1));
        t.line("-- after %d s in total:", wait);
        assertThat(call("alice").getResponse().getStatus()).isEqualTo(200);
        assertThat(call("alice").getResponse().getStatus()).isEqualTo(429);
    }

    @Test
    @Order(4)
    void healthIsNotLimited() throws Exception {
        t.line("== /health is outside /api/, so the filter skips it");
        for (int i = 0; i < 30; i++) {
            assertThat(mvc.perform(get("/health")).andReturn().getResponse().getStatus()).isEqualTo(200);
        }
        t.line("30 requests to /health -> all 200, X-RateLimit-Limit header present: %s",
                mvc.perform(get("/health")).andReturn().getResponse().containsHeader("X-RateLimit-Limit"));
    }
}