Spaces:
Sleeping
Sleeping
File size: 5,605 Bytes
ba5110e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 | """
Test cases for Rate Limiting module.
Tests GPT-OSS limits and Wolfram monthly limits.
"""
import pytest
import time
from backend.utils.rate_limit import (
RateLimitTracker,
SessionRateLimiter,
WolframRateLimiter,
QueryCache,
RATE_LIMITS,
WOLFRAM_MONTHLY_LIMIT,
)
class TestRateLimitTracker:
"""Test suite for session rate limit tracking."""
def test_initial_state(self):
"""TC-RL-001: Initial tracker should allow requests."""
tracker = RateLimitTracker()
can_proceed, msg = tracker.can_make_request()
assert can_proceed is True
assert msg == ""
def test_record_usage(self):
"""TC-RL-002: Recording usage should increment counters."""
tracker = RateLimitTracker()
tracker.record_usage(100)
assert tracker.requests_this_minute == 1
assert tracker.tokens_this_minute == 100
def test_rpm_limit(self):
"""TC-RL-003: Should block after exceeding RPM limit."""
tracker = RateLimitTracker()
# Simulate 30 requests
for _ in range(30):
tracker.record_usage(10)
can_proceed, msg = tracker.can_make_request()
assert can_proceed is False
assert "Rate limit" in msg or "wait" in msg.lower()
def test_token_limit(self):
"""TC-RL-004: Should block after exceeding TPM limit."""
tracker = RateLimitTracker()
# Record close to 8000 tokens
tracker.tokens_this_minute = 7500
can_proceed, msg = tracker.can_make_request(estimated_tokens=1000)
assert can_proceed is False
assert "Token" in msg or "limit" in msg.lower()
def test_daily_limit(self):
"""TC-RL-005: Should block after exceeding daily requests."""
tracker = RateLimitTracker()
tracker.requests_today = RATE_LIMITS["rpd"]
can_proceed, msg = tracker.can_make_request()
assert can_proceed is False
assert "Daily" in msg or "tomorrow" in msg.lower()
class TestSessionRateLimiter:
"""Test suite for multi-session rate limiting."""
def test_separate_sessions(self):
"""TC-RL-006: Different sessions should have independent limits."""
limiter = SessionRateLimiter()
# Record usage for session A
limiter.record("session_a", 100)
# Session B should still be clean
tracker_b = limiter.get_tracker("session_b")
assert tracker_b.requests_this_minute == 0
def test_session_persistence(self):
"""TC-RL-007: Same session should accumulate usage."""
limiter = SessionRateLimiter()
limiter.record("session_x", 50)
limiter.record("session_x", 50)
tracker = limiter.get_tracker("session_x")
assert tracker.requests_this_minute == 2
assert tracker.tokens_this_minute == 100
class TestWolframRateLimiter:
"""Test suite for Wolfram Alpha monthly rate limiting."""
def test_initial_usage(self):
"""TC-RL-008: Initial usage should be 0 or existing value."""
limiter = WolframRateLimiter(cache_dir=".test_caches/wolfram_cache")
status = limiter.get_status()
assert status["limit"] == WOLFRAM_MONTHLY_LIMIT
assert isinstance(status["used"], int)
assert isinstance(status["remaining"], int)
def test_can_make_request_initially(self):
"""TC-RL-009: Should allow requests when under limit."""
limiter = WolframRateLimiter(cache_dir=".test_caches/wolfram_cache_2")
can_proceed, msg, remaining = limiter.can_make_request()
assert can_proceed is True
def test_record_increments_usage(self):
"""TC-RL-010: Recording should increment usage counter."""
limiter = WolframRateLimiter(cache_dir=".test_caches/wolfram_cache_3")
initial = limiter.get_usage()
limiter.record_usage()
after = limiter.get_usage()
assert after == initial + 1
def test_month_key_format(self):
"""TC-RL-011: Month key should be in correct format."""
limiter = WolframRateLimiter()
key = limiter._get_month_key()
assert key.startswith("wolfram_usage_")
assert "2025" in key # Current year
class TestQueryCache:
"""Test suite for query caching."""
def test_cache_miss(self):
"""TC-RL-012: Non-existent query should return None."""
cache = QueryCache(cache_dir=".test_caches/cache_1")
result = cache.get("nonexistent_query_12345")
assert result is None
def test_cache_set_and_get(self):
"""TC-RL-013: Cached query should be retrievable."""
cache = QueryCache(cache_dir=".test_caches/cache_2")
cache.set("test_query", "test_response", context="test")
result = cache.get("test_query", context="test")
assert result == "test_response"
def test_cache_context_separation(self):
"""TC-RL-014: Different contexts should have separate caches."""
cache = QueryCache(cache_dir=".test_caches/cache_3")
cache.set("query", "response_a", context="context_a")
cache.set("query", "response_b", context="context_b")
assert cache.get("query", context="context_a") == "response_a"
assert cache.get("query", context="context_b") == "response_b"
def test_cache_clear(self):
"""TC-RL-015: Clear should remove all cached entries."""
cache = QueryCache(cache_dir=".test_caches/cache_4")
cache.set("key1", "value1")
cache.clear()
assert cache.get("key1") is None
|