update

TabbyML · Dec 19, 2024 · 2c32784 · 2c32784
1 parent 684b985
commit 2c32784
Showing 1 changed file with 8 additions and 2 deletions.
diff --git 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);
     }
 }