use axum::{routing::get, Router};
use http::{Request, StatusCode};
use std::sync::Arc;
use touch_ratelimit::{
adapters::axum::axum_rate_limit_layer, bucket::token_bucket::TokenBucket,
storage::in_memory::InMemoryStore,
};
use tower::ServiceExt;
#[tokio::test]
async fn axum_rate_limit_works() {
let store = Arc::new(InMemoryStore::new(|| TokenBucket::new(1.0, 0.0)));
let app = Router::new()
.route("/", get(|| async { "ok" }))
.layer(axum_rate_limit_layer(store));
let response = app
.clone()
.oneshot(
Request::builder()
.uri("/")
.header("x-forwarded-for", "1.2.3.4")
.body(axum::body::Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let response = app
.oneshot(
Request::builder()
.uri("/")
.header("x-forwarded-for", "1.2.3.4")
.body(axum::body::Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
}