From 6977494c53b729b5dbb62eee21cd9c09a33c8eba Mon Sep 17 00:00:00 2001 From: DazedAnon Date: Mon, 22 Jun 2026 05:42:12 -0500 Subject: [PATCH] Rate Limiter Unit Tests --- tests/test_rate_limiter.py | 228 +++++++++++++++++++++++++++++++++++++ 1 file changed, 228 insertions(+) create mode 100644 tests/test_rate_limiter.py diff --git a/tests/test_rate_limiter.py b/tests/test_rate_limiter.py new file mode 100644 index 0000000..d5d17e2 --- /dev/null +++ b/tests/test_rate_limiter.py @@ -0,0 +1,228 @@ +import types +import unittest +from unittest import mock + +from openai import APIStatusError, RateLimitError + +import util.translation as T +from util.translation import AdaptiveLimiter, _header_float, _header_int + + +class FakeClock: + """Deterministic stand-in for the ``time`` module. + + ``sleep`` advances the clock instead of waiting, so pacing/backoff paths + resolve instantly and predictably. + """ + + def __init__(self, start=1000.0): + self.t = start + + def monotonic(self): + return self.t + + def time(self): + return self.t + + def sleep(self, seconds): + self.t += seconds + + +def make_exc(cls, **attrs): + """Build an openai exception without running its real constructor.""" + err = cls.__new__(cls) + for key, value in attrs.items(): + setattr(err, key, value) + return err + + +class FakeRaw: + def __init__(self, headers, parsed): + self._headers = headers + self._parsed = parsed + + @property + def headers(self): + return self._headers + + def parse(self): + return self._parsed + + +class FakeCreate: + """Callable that replays a scripted sequence of returns/raises.""" + + def __init__(self, behaviors): + self.behaviors = behaviors + self.calls = [] + + def __call__(self, **params): + self.calls.append(params) + behavior = self.behaviors[len(self.calls) - 1] + if isinstance(behavior, Exception): + raise behavior + return behavior + + +def fake_openai(create): + completions = types.SimpleNamespace( + with_raw_response=types.SimpleNamespace(create=create) + ) + return types.SimpleNamespace(chat=types.SimpleNamespace(completions=completions)) + + +class HeaderHelperTests(unittest.TestCase): + def test_header_int_first_present_and_floored(self): + headers = {"b": "125.9"} + self.assertEqual(_header_int(headers, "a", "b"), 125) + + def test_header_int_missing_blank_and_garbage(self): + self.assertIsNone(_header_int({}, "a")) + self.assertIsNone(_header_int({"a": " "}, "a")) + self.assertIsNone(_header_int({"a": "abc"}, "a")) + + def test_header_float_keeps_fraction(self): + self.assertAlmostEqual(_header_float({"a": "0.42"}, "a"), 0.42) + + +class AdaptiveLimiterUpdateTests(unittest.TestCase): + def _limiter(self): + # Seeded at 0.5 RPS / 50k TPM, matching the production default. + return AdaptiveLimiter(0.5, 50000, 4000) + + def test_canonical_headers_correct_pace_and_budget(self): + lim = self._limiter() + lim.update({ + "x-ratelimit-limit-requests": "125", + "x-ratelimit-limit-tokens": "500000", + "x-ratelimit-remaining-tokens": "498000", + }) + self.assertAlmostEqual(lim.req_per_sec, 125 / 60.0) + self.assertAlmostEqual(lim.min_interval, 60.0 / 125) + self.assertEqual(lim.tok_per_min, 500000) + self.assertEqual(lim.tokens_remaining, 498000) + + def test_legacy_minute_spelling_still_works(self): + lim = self._limiter() + lim.update({ + "x-ratelimit-limit-req-minute": "25", + "x-ratelimit-limit-tokens-minute": "60000", + "x-ratelimit-remaining-tokens-minute": "59000", + }) + self.assertAlmostEqual(lim.req_per_sec, 25 / 60.0) + self.assertEqual(lim.tok_per_min, 60000) + self.assertEqual(lim.tokens_remaining, 59000) + + def test_true_per_second_header_used_undivided(self): + lim = self._limiter() + lim.update({"x-ratelimit-limit-req-second": "2"}) + self.assertAlmostEqual(lim.req_per_sec, 2.0) + + def test_empty_or_missing_headers_leave_seed_untouched(self): + for headers in (None, {}): + lim = self._limiter() + lim.update(headers) + self.assertAlmostEqual(lim.req_per_sec, 0.5) + self.assertEqual(lim.tok_per_min, 50000) + + def test_garbage_values_ignored(self): + lim = self._limiter() + lim.update({"x-ratelimit-limit-requests": "abc"}) + self.assertAlmostEqual(lim.req_per_sec, 0.5) + + def test_exhausted_request_budget_holds_off(self): + clock = FakeClock() + with mock.patch.object(T, "time", clock): + lim = self._limiter() + lim.update({"x-ratelimit-remaining-requests": "0"}) + self.assertGreaterEqual(lim.next_request_at, clock.monotonic() + 60.0) + + +class AdaptiveLimiterAcquireTests(unittest.TestCase): + def test_requests_spaced_by_min_interval(self): + clock = FakeClock() + with mock.patch.object(T, "time", clock): + lim = AdaptiveLimiter(2.0, 10_000_000, 4000) # min_interval = 0.5s + lim.acquire(10) # first clears immediately + self.assertEqual(clock.t, 1000.0) + lim.acquire(10) # second waits one interval + self.assertAlmostEqual(clock.t, 1000.5) + self.assertEqual(lim.tokens_remaining, 10_000_000 - 20) + + def test_token_gate_waits_for_window_reset(self): + clock = FakeClock() + with mock.patch.object(T, "time", clock): + lim = AdaptiveLimiter(1000.0, 10000, 2000) + lim.tokens_remaining = 2500 # near depletion: 2500 - 1000 <= 2000 + original_reset = lim.window_reset + lim.acquire(1000) + # It blocked until the minute window rolled over. + self.assertGreaterEqual(clock.t, original_reset) + # After the reset the budget refills, then this call is charged. + self.assertEqual(lim.tokens_remaining, 10000 - 1000) + + +class CallMistralTests(unittest.TestCase): + def setUp(self): + self._saved_limiter = T._mistral_limiter + T._mistral_limiter = None # force a fresh limiter from the seeds + self.clock = FakeClock() + self._patches = [ + mock.patch.object(T, "time", self.clock), + mock.patch.dict(T.os.environ, {}, clear=False), + ] + for patch in self._patches: + patch.start() + + def tearDown(self): + for patch in self._patches: + patch.stop() + T._mistral_limiter = self._saved_limiter + + def _params(self, **extra): + base = {"messages": [{"content": "ใƒ†ใ‚นใƒˆ"}], "max_tokens": 100} + base.update(extra) + return base + + def test_success_returns_response_and_syncs_limiter(self): + sentinel = object() + create = FakeCreate([ + FakeRaw({"x-ratelimit-limit-requests": "125"}, sentinel), + ]) + with mock.patch.object(T, "openai", fake_openai(create)): + result = T.callMistral(self._params()) + self.assertIs(result, sentinel) + self.assertEqual(len(create.calls), 1) + self.assertAlmostEqual(T._get_mistral_limiter().req_per_sec, 125 / 60.0) + + def test_rate_limit_then_success_retries(self): + sentinel = object() + rate_limited = make_exc( + RateLimitError, + response=types.SimpleNamespace(headers={"retry-after": "0.01"}), + ) + create = FakeCreate([ + rate_limited, + FakeRaw({"x-ratelimit-limit-requests": "60"}, sentinel), + ]) + with mock.patch.object(T, "openai", fake_openai(create)): + result = T.callMistral(self._params()) + self.assertIs(result, sentinel) + self.assertEqual(len(create.calls), 2) + + def test_json_schema_downgraded_to_json_object(self): + sentinel = object() + bad_format = make_exc(APIStatusError, status_code=400) + create = FakeCreate([ + bad_format, + FakeRaw({}, sentinel), + ]) + params = self._params(response_format={"type": "json_schema"}) + with mock.patch.object(T, "openai", fake_openai(create)): + result = T.callMistral(params) + self.assertIs(result, sentinel) + self.assertEqual(create.calls[1]["response_format"], {"type": "json_object"}) + + +if __name__ == "__main__": + unittest.main()