mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
ea1951508b
## Description Extracts token-bucket refill, consume, wait-time, and stale-bucket selection formulas into a pure rate-limit policy module while preserving the async `TokenBucketRateLimiter` adapter for locks and mutable bucket storage. Closes # ## Type of Change - [ ] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [x] Code refactoring (no functional changes) ## Changes Made - Added `headroom.proxy.rate_limit_policy` for pure token-bucket calculations. - Updated `TokenBucketRateLimiter` to delegate refill, consume, and stale-key selection to the extracted policy. - Added direct tests for the rate-limit policy boundary. - Kept the LiteLLM callback compatibility shim required for repo-wide type checking on fresh branches. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text python -m pytest tests/test_rate_limit_policy.py tests/test_proxy_healthchecks.py tests/test_litellm_callback.py tests/test_compress_api.py::TestLiteLLMCallback -q 27 passed in 17.82s python -m ruff check . All checks passed! python -m ruff format --check . 1095 files already formatted python -m mypy headroom --ignore-missing-imports Success: no issues found in 409 source files gitleaks protect --staged --no-banner --redact no leaks found ``` ## Real Behavior Proof - Environment: Windows, Python 3.13.13, local worktree based on `headroomlabs/main`. - Exact command / steps: ran focused rate-limit policy tests, proxy health checks, LiteLLM callback tests, Ruff lint/format checks, mypy over `headroom`, and staged gitleaks scan. - Observed result: all local checks passed; staged secret scan found no leaks. - Not tested: full CI matrix and deployment flows; those are covered by GitHub Actions. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable ## Screenshots (if applicable) N/A ## Additional Notes Documentation and changelog updates are not applicable for this internal refactor. GitHub reported existing Dependabot alerts on the default branch during push; this PR does not change dependencies, and the staged secret scan is clean.
65 lines
1.5 KiB
Python
65 lines
1.5 KiB
Python
"""Tests for pure token-bucket rate-limit policy helpers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from headroom.proxy.rate_limit_policy import (
|
|
consume_from_bucket,
|
|
refilled_tokens,
|
|
stale_bucket_keys,
|
|
)
|
|
|
|
|
|
def test_refilled_tokens_caps_at_bucket_rate() -> None:
|
|
assert (
|
|
refilled_tokens(
|
|
current_tokens=9,
|
|
last_update=0,
|
|
now=120,
|
|
rate_per_minute=10,
|
|
)
|
|
== 10
|
|
)
|
|
|
|
|
|
def test_refilled_tokens_ignores_negative_elapsed_time() -> None:
|
|
assert (
|
|
refilled_tokens(
|
|
current_tokens=3,
|
|
last_update=10,
|
|
now=5,
|
|
rate_per_minute=60,
|
|
)
|
|
== 3
|
|
)
|
|
|
|
|
|
def test_consume_from_bucket_allows_and_debits_available_tokens() -> None:
|
|
allowed, remaining, wait_seconds = consume_from_bucket(
|
|
available_tokens=5,
|
|
requested_tokens=2,
|
|
rate_per_minute=60,
|
|
)
|
|
|
|
assert allowed is True
|
|
assert remaining == 3
|
|
assert wait_seconds == 0
|
|
|
|
|
|
def test_consume_from_bucket_denies_and_reports_wait_time() -> None:
|
|
allowed, remaining, wait_seconds = consume_from_bucket(
|
|
available_tokens=0.5,
|
|
requested_tokens=1,
|
|
rate_per_minute=60,
|
|
)
|
|
|
|
assert allowed is False
|
|
assert remaining == 0.5
|
|
assert wait_seconds == 0.5
|
|
|
|
|
|
def test_stale_bucket_keys_returns_only_old_buckets() -> None:
|
|
assert stale_bucket_keys(
|
|
{"fresh": 950, "edge": 400, "stale": 399},
|
|
now=1000,
|
|
stale_after_seconds=600,
|
|
) == ["stale"]
|