fastapi-traffic/tests/test_backends.py

"""Tests for rate limit backends."""

from __future__ import annotations

import asyncio
from typing import AsyncGenerator

import pytest

from fastapi_traffic.backends.memory import MemoryBackend
from fastapi_traffic.backends.sqlite import SQLiteBackend


class TestMemoryBackend:
    """Tests for MemoryBackend."""

    @pytest.fixture
    async def backend(self) -> AsyncGenerator[MemoryBackend, None]:
        """Create a memory backend for testing."""
        backend = MemoryBackend(max_size=100, cleanup_interval=1.0)
        yield backend
        await backend.close()

    async def test_set_and_get(self, backend: MemoryBackend) -> None:
        """Test basic set and get operations."""
        await backend.set("test_key", {"count": 5}, ttl=60.0)
        result = await backend.get("test_key")
        assert result is not None
        assert result["count"] == 5

    async def test_get_nonexistent(self, backend: MemoryBackend) -> None:
        """Test getting a nonexistent key."""
        result = await backend.get("nonexistent")
        assert result is None

    async def test_delete(self, backend: MemoryBackend) -> None:
        """Test delete operation."""
        await backend.set("test_key", {"count": 5}, ttl=60.0)
        await backend.delete("test_key")
        result = await backend.get("test_key")
        assert result is None

    async def test_exists(self, backend: MemoryBackend) -> None:
        """Test exists operation."""
        assert not await backend.exists("test_key")
        await backend.set("test_key", {"count": 5}, ttl=60.0)
        assert await backend.exists("test_key")

    async def test_increment(self, backend: MemoryBackend) -> None:
        """Test increment operation."""
        await backend.set("test_key", {"count": 5}, ttl=60.0)
        result = await backend.increment("test_key", 3)
        assert result == 8

    async def test_clear(self, backend: MemoryBackend) -> None:
        """Test clear operation."""
        await backend.set("key1", {"count": 1}, ttl=60.0)
        await backend.set("key2", {"count": 2}, ttl=60.0)
        await backend.clear()
        assert not await backend.exists("key1")
        assert not await backend.exists("key2")

    async def test_ttl_expiration(self, backend: MemoryBackend) -> None:
        """Test that entries expire after TTL."""
        await backend.set("test_key", {"count": 5}, ttl=0.1)
        await asyncio.sleep(0.2)
        result = await backend.get("test_key")
        assert result is None

    async def test_lru_eviction(self) -> None:
        """Test LRU eviction when max size is reached."""
        backend = MemoryBackend(max_size=3)
        try:
            await backend.set("key1", {"v": 1}, ttl=60.0)
            await backend.set("key2", {"v": 2}, ttl=60.0)
            await backend.set("key3", {"v": 3}, ttl=60.0)
            await backend.set("key4", {"v": 4}, ttl=60.0)

            assert not await backend.exists("key1")
            assert await backend.exists("key2")
            assert await backend.exists("key3")
            assert await backend.exists("key4")
        finally:
            await backend.close()


class TestSQLiteBackend:
    """Tests for SQLiteBackend."""

    @pytest.fixture
    async def backend(self) -> AsyncGenerator[SQLiteBackend, None]:
        """Create an in-memory SQLite backend for testing."""
        backend = SQLiteBackend(":memory:", cleanup_interval=1.0)
        await backend.initialize()
        yield backend
        await backend.close()

    async def test_set_and_get(self, backend: SQLiteBackend) -> None:
        """Test basic set and get operations."""
        await backend.set("test_key", {"count": 5}, ttl=60.0)
        result = await backend.get("test_key")
        assert result is not None
        assert result["count"] == 5

    async def test_get_nonexistent(self, backend: SQLiteBackend) -> None:
        """Test getting a nonexistent key."""
        result = await backend.get("nonexistent")
        assert result is None

    async def test_delete(self, backend: SQLiteBackend) -> None:
        """Test delete operation."""
        await backend.set("test_key", {"count": 5}, ttl=60.0)
        await backend.delete("test_key")
        result = await backend.get("test_key")
        assert result is None

    async def test_exists(self, backend: SQLiteBackend) -> None:
        """Test exists operation."""
        assert not await backend.exists("test_key")
        await backend.set("test_key", {"count": 5}, ttl=60.0)
        assert await backend.exists("test_key")

    async def test_increment(self, backend: SQLiteBackend) -> None:
        """Test increment operation."""
        await backend.set("test_key", {"count": 5}, ttl=60.0)
        result = await backend.increment("test_key", 3)
        assert result == 8

    async def test_clear(self, backend: SQLiteBackend) -> None:
        """Test clear operation."""
        await backend.set("key1", {"count": 1}, ttl=60.0)
        await backend.set("key2", {"count": 2}, ttl=60.0)
        await backend.clear()
        assert not await backend.exists("key1")
        assert not await backend.exists("key2")

    async def test_get_stats(self, backend: SQLiteBackend) -> None:
        """Test get_stats operation."""
        await backend.set("key1", {"count": 1}, ttl=60.0)
        await backend.set("key2", {"count": 2}, ttl=60.0)
        stats = await backend.get_stats()
        assert stats["total_entries"] == 2
        assert stats["active_entries"] == 2