#![cfg(feature = "testing")]
use ultimo::middleware::builtin::{rate_limiter, RateLimitKey, RateLimiter};
use ultimo::testing::TestClient;
use ultimo::{Context, Ultimo};
#[tokio::test]
async fn rate_limiter_allows_within_limit() {
let mut app = Ultimo::new_without_defaults();
app.use_middleware(RateLimiter::new(5, 60).key(RateLimitKey::Global).build());
app.get("/", |ctx: Context| async move { ctx.text("ok").await });
let client = TestClient::new(app);
for _ in 0..5 {
let res = client.get("/").send().await;
assert_eq!(res.status(), 200);
}
}
#[tokio::test]
async fn rate_limiter_returns_429_when_exceeded() {
let mut app = Ultimo::new_without_defaults();
app.use_middleware(RateLimiter::new(3, 60).key(RateLimitKey::Global).build());
app.get("/", |ctx: Context| async move { ctx.text("ok").await });
let client = TestClient::new(app);
for _ in 0..3 {
let res = client.get("/").send().await;
assert_eq!(res.status(), 200);
}
let res = client.get("/").send().await;
assert_eq!(res.status(), 429);
assert!(res.header("retry-after").is_some());
}
#[tokio::test]
async fn rate_limiter_default_convenience_function() {
let mut app = Ultimo::new_without_defaults();
app.use_middleware(rate_limiter());
app.get("/", |ctx: Context| async move { ctx.text("ok").await });
let client = TestClient::new(app);
let res = client.get("/").send().await;
assert_eq!(res.status(), 200);
}
#[tokio::test]
async fn rate_limiter_per_header_key() {
let mut app = Ultimo::new_without_defaults();
app.use_middleware(
RateLimiter::new(2, 60)
.key(RateLimitKey::Header("X-API-Key".into()))
.build(),
);
app.get("/", |ctx: Context| async move { ctx.text("ok").await });
let client = TestClient::new(app);
let res = client.get("/").header("X-API-Key", "client-a").send().await;
assert_eq!(res.status(), 200);
let res = client.get("/").header("X-API-Key", "client-a").send().await;
assert_eq!(res.status(), 200);
let res = client.get("/").header("X-API-Key", "client-a").send().await;
assert_eq!(res.status(), 429);
let res = client.get("/").header("X-API-Key", "client-b").send().await;
assert_eq!(res.status(), 200);
}