From 2c32784ee7413ce02113767c6550d01bdbb17500 Mon Sep 17 00:00:00 2001 From: Meng Zhang Date: Thu, 19 Dec 2024 09:35:06 +0800 Subject: [PATCH] update --- ee/tabby-webserver/src/rate_limit.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/ee/tabby-webserver/src/rate_limit.rs b/ee/tabby-webserver/src/rate_limit.rs index ddbc0f5b0352..f313a98bb5e7 100644 --- a/ee/tabby-webserver/src/rate_limit.rs +++ b/ee/tabby-webserver/src/rate_limit.rs @@ -55,12 +55,18 @@ mod tests { let user_id = "test_user"; let rate_limiter = UserRateLimiter::default(); + let uri: axum::http::Uri = "/v1/completions".parse().unwrap(); + let healthcheck_uri: axum::http::Uri = "/v1/health".parse().unwrap(); + // Test that the first `USER_REQUEST_LIMIT_PER_MINUTE` requests are allowed for _ in 0..USER_REQUEST_LIMIT_PER_MINUTE { - assert!(rate_limiter.is_allowed(user_id).await); + assert!(rate_limiter.is_allowed(&uri, user_id).await); } // Test that the 201st request is not allowed - assert!(!rate_limiter.is_allowed(user_id).await); + assert!(!rate_limiter.is_allowed(&uri, user_id).await); + + // Test that health check requests are not limited + assert!(rate_limiter.is_allowed(&healthcheck_uri, user_id).await); } }