Skip to content

Commit df031b4

Browse files
committed
pr comment updates, add tests, move sempahore
1 parent 915de1a commit df031b4

File tree

1 file changed

+89
-0
lines changed

1 file changed

+89
-0
lines changed

src/rate_limit.rs

+89
Original file line numberDiff line numberDiff line change
@@ -599,6 +599,95 @@ mod tests {
599599
assert!(c4.is_ok());
600600
}
601601

602+
#[tokio::test]
603+
async fn test_global_limits_with_multiple_ips() {
604+
let user_1 = IpAddr::from_str("127.0.0.1").unwrap();
605+
let user_2 = IpAddr::from_str("127.0.0.2").unwrap();
606+
let user_3 = IpAddr::from_str("127.0.0.3").unwrap();
607+
608+
let rate_limiter = Arc::new(InMemoryRateLimit::new(4, 3));
609+
610+
let ticket_1_1 = rate_limiter.clone().try_acquire(user_1).unwrap();
611+
let ticket_1_2 = rate_limiter.clone().try_acquire(user_1).unwrap();
612+
613+
let ticket_2_1 = rate_limiter.clone().try_acquire(user_2).unwrap();
614+
let ticket_2_2 = rate_limiter.clone().try_acquire(user_2).unwrap();
615+
616+
assert_eq!(
617+
rate_limiter.inner.lock().unwrap().semaphore.available_permits(),
618+
0
619+
);
620+
621+
// Try user_3 - should fail due to global limit
622+
let result = rate_limiter.clone().try_acquire(user_3);
623+
assert!(result.is_err());
624+
assert_eq!(
625+
result.err().unwrap().to_string(),
626+
"Rate Limit Reached: Global limit"
627+
);
628+
629+
drop(ticket_1_1);
630+
631+
let ticket_3_1 = rate_limiter.clone().try_acquire(user_3).unwrap();
632+
633+
drop(ticket_1_2);
634+
drop(ticket_2_1);
635+
drop(ticket_2_2);
636+
drop(ticket_3_1);
637+
638+
assert_eq!(
639+
rate_limiter.inner.lock().unwrap().semaphore.available_permits(),
640+
4
641+
);
642+
assert_eq!(
643+
rate_limiter.inner.lock().unwrap().active_connections.len(),
644+
0
645+
);
646+
}
647+
648+
#[tokio::test]
649+
async fn test_per_ip_limits_remain_enforced() {
650+
let user_1 = IpAddr::from_str("127.0.0.1").unwrap();
651+
let user_2 = IpAddr::from_str("127.0.0.2").unwrap();
652+
653+
let rate_limiter = Arc::new(InMemoryRateLimit::new(5, 2));
654+
655+
let ticket_1_1 = rate_limiter.clone().try_acquire(user_1).unwrap();
656+
let ticket_1_2 = rate_limiter.clone().try_acquire(user_1).unwrap();
657+
658+
let result = rate_limiter.clone().try_acquire(user_1);
659+
assert!(result.is_err());
660+
assert_eq!(
661+
result.err().unwrap().to_string(),
662+
"Rate Limit Reached: IP limit exceeded"
663+
);
664+
665+
let ticket_2_1 = rate_limiter.clone().try_acquire(user_2).unwrap();
666+
drop(ticket_1_1);
667+
668+
let ticket_1_3 = rate_limiter.clone().try_acquire(user_1).unwrap();
669+
670+
let result = rate_limiter.clone().try_acquire(user_1);
671+
assert!(result.is_err());
672+
assert_eq!(
673+
result.err().unwrap().to_string(),
674+
"Rate Limit Reached: IP limit exceeded"
675+
);
676+
677+
drop(ticket_1_2);
678+
drop(ticket_1_3);
679+
drop(ticket_2_1);
680+
681+
assert_eq!(
682+
rate_limiter.inner.lock().unwrap().semaphore.available_permits(),
683+
5
684+
);
685+
assert_eq!(
686+
rate_limiter.inner.lock().unwrap().active_connections.len(),
687+
0
688+
);
689+
}
690+
602691
#[tokio::test]
603692
#[cfg(all(feature = "integration", test))]
604693
async fn test_instance_tracking_and_cleanup() {

0 commit comments

Comments
 (0)