Aether/tests/plugins/test_token_bucket.py

from __future__ import annotations

from datetime import datetime, timezone

import pytest

import src.plugins.rate_limit.token_bucket as token_bucket_module
from src.plugins.rate_limit.token_bucket import RedisTokenBucketBackend, TokenBucketStrategy


@pytest.mark.asyncio
async def test_token_bucket_cleans_up_expired_buckets(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("RATE_LIMIT_BACKEND", "memory")
    strategy = TokenBucketStrategy()
    strategy.configure({"bucket_expiry": 1, "cleanup_interval": 0})

    await strategy.check_limit("api_key:stale")
    strategy.buckets["api_key:stale"].last_access_time -= 3600

    await strategy.check_limit("api_key:fresh")

    assert "api_key:stale" not in strategy.buckets
    assert "api_key:fresh" in strategy.buckets


@pytest.mark.asyncio
async def test_token_bucket_reconfigures_existing_bucket_when_rate_limit_changes(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setenv("RATE_LIMIT_BACKEND", "memory")
    strategy = TokenBucketStrategy()

    await strategy.check_limit("user:42", rate_limit=120)
    bucket = strategy.buckets["user:42"]
    bucket.tokens = 90

    await strategy.check_limit("user:42", rate_limit=30)

    updated_bucket = strategy.buckets["user:42"]
    assert updated_bucket.capacity == 30
    assert updated_bucket.refill_rate == 0.5
    assert updated_bucket.tokens <= 30


@pytest.mark.asyncio
async def test_token_bucket_treats_non_positive_dynamic_rate_limit_as_unlimited(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setenv("RATE_LIMIT_BACKEND", "memory")
    strategy = TokenBucketStrategy()

    result = await strategy.check_limit("public_ip:test", rate_limit=0)
    consumed = await strategy.consume("public_ip:test", amount=1, rate_limit=0)

    assert result.allowed is True
    assert consumed is True
    assert "public_ip:test" not in strategy.buckets


class _FakeRedisClient:
    async def hmget(self, _key: str, *_fields: str) -> list[None]:
        return [None, None]

    def register_script(self, _script: str):  # type: ignore[no-untyped-def]
        async def _runner(*args, **kwargs):  # type: ignore[no-untyped-def]
            return [1, 0, 0]

        return _runner


@pytest.mark.asyncio
async def test_token_bucket_retries_redis_backend_probe_after_initial_miss(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setenv("RATE_LIMIT_BACKEND", "auto")
    strategy = TokenBucketStrategy()
    strategy._redis_retry_interval = 0

    fake_redis = _FakeRedisClient()
    calls = {"count": 0}

    def _fake_get_redis_client_sync():  # type: ignore[no-untyped-def]
        calls["count"] += 1
        if calls["count"] == 1:
            return None
        return fake_redis

    monkeypatch.setattr(
        token_bucket_module,
        "get_redis_client_sync",
        _fake_get_redis_client_sync,
    )

    await strategy.check_limit("public_ip:first")
    assert strategy._redis_backend is None

    await strategy.check_limit("public_ip:second")
    assert strategy._redis_backend is not None
    assert calls["count"] == 2


@pytest.mark.asyncio
async def test_redis_token_bucket_missing_bucket_reports_reset_now() -> None:
    backend = RedisTokenBucketBackend(_FakeRedisClient())

    result = await backend.peek("public_ip:test", capacity=60, refill_rate=1.0, amount=1)

    assert result.allowed is True
    assert result.remaining == 60
    assert result.reset_at is not None
    assert abs((result.reset_at - datetime.now(timezone.utc)).total_seconds()) < 2