r2r/tests/unit/test_limits.py

import uuid
from datetime import datetime, timedelta, timezone
from uuid import UUID

import pytest

from core.base import LimitSettings
from core.database.postgres import PostgresLimitsHandler
from shared.abstractions import User


@pytest.mark.asyncio
async def test_log_request_and_count(limits_handler):
    """
    Test that when we log requests, the count increments, and rate-limits are enforced.
    Route-specific test using the /v3/retrieval/search endpoint limits.
    """
    # Clear existing logs first
    clear_query = f"DELETE FROM {limits_handler._get_table_name(PostgresLimitsHandler.TABLE_NAME)}"
    await limits_handler.connection_manager.execute_query(clear_query)

    user_id = uuid.uuid4()
    route = "/v3/retrieval/search"  # Using actual route from config
    test_user = User(
        id=user_id,
        email="test@example.com",
        is_active=True,
        is_verified=True,
        is_superuser=False,
        limits_overrides=None,
    )

    # Set route limit to match config: 5 requests per minute
    old_route_limits = limits_handler.config.route_limits
    new_route_limits = {
        route: LimitSettings(route_per_min=5, monthly_limit=10)
    }
    limits_handler.config.route_limits = new_route_limits

    print(f"\nTesting with route limits: {new_route_limits}")
    print(f"Route settings: {limits_handler.config.route_limits[route]}")

    try:
        # Initial check should pass (no requests yet)
        await limits_handler.check_limits(test_user, route)
        print("Initial check passed (no requests)")

        # Log 5 requests (exactly at limit)
        for i in range(5):
            await limits_handler.log_request(user_id, route)
            now = datetime.now(timezone.utc)
            one_min_ago = now - timedelta(minutes=1)
            route_count = await limits_handler._count_requests(
                user_id, route, one_min_ago
            )
            print(f"Route count after request {i+1}: {route_count}")

            # This should pass for all 5 requests
            await limits_handler.check_limits(test_user, route)
            print(f"Check limits passed after request {i+1}")

        # Log the 6th request (over limit)
        await limits_handler.log_request(user_id, route)
        route_count = await limits_handler._count_requests(
            user_id, route, one_min_ago
        )
        print(f"Route count after request 6: {route_count}")

        # This check should fail as we've exceeded route_per_min=5
        with pytest.raises(
            ValueError, match="Per-route per-minute rate limit exceeded"
        ):
            await limits_handler.check_limits(test_user, route)

    finally:
        limits_handler.config.route_limits = old_route_limits


@pytest.mark.asyncio
async def test_global_limit(limits_handler):
    """
    Test global limit using the configured limit of 10 requests per minute
    """
    # Clear existing logs
    clear_query = f"DELETE FROM {limits_handler._get_table_name(PostgresLimitsHandler.TABLE_NAME)}"
    await limits_handler.connection_manager.execute_query(clear_query)

    user_id = uuid.uuid4()
    route = "/global-test"
    test_user = User(
        id=user_id,
        email="globaltest@example.com",
        is_active=True,
        is_verified=True,
        is_superuser=False,
        limits_overrides=None,
    )

    # Set global limit to match config: 10 requests per minute
    old_limits = limits_handler.config.limits
    limits_handler.config.limits = LimitSettings(
        global_per_min=10, monthly_limit=20
    )

    try:
        # Initial check should pass (no requests)
        await limits_handler.check_limits(test_user, route)
        print("Initial global check passed (no requests)")

        # Log 10 requests (hits the limit)
        for i in range(11):
            await limits_handler.log_request(user_id, route)

        # Debug counts
        now = datetime.now(timezone.utc)
        one_min_ago = now - timedelta(minutes=1)
        global_count = await limits_handler._count_requests(
            user_id, None, one_min_ago
        )
        print(f"Global count after 10 requests: {global_count}")

        # This should fail as we've hit global_per_min=10
        with pytest.raises(
            ValueError, match="Global per-minute rate limit exceeded"
        ):
            await limits_handler.check_limits(test_user, route)

    finally:
        limits_handler.config.limits = old_limits


@pytest.mark.asyncio
async def test_monthly_limit(limits_handler):
    """
    Test monthly limit using the configured limit of 20 requests per month
    """
    # Clear existing logs
    clear_query = f"DELETE FROM {limits_handler._get_table_name(PostgresLimitsHandler.TABLE_NAME)}"
    await limits_handler.connection_manager.execute_query(clear_query)

    user_id = uuid.uuid4()
    route = "/monthly-test"
    test_user = User(
        id=user_id,
        email="monthly@example.com",
        is_active=True,
        is_verified=True,
        is_superuser=False,
        limits_overrides=None,
    )

    old_limits = limits_handler.config.limits
    limits_handler.config.limits = LimitSettings(monthly_limit=20)

    try:
        # Initial check should pass (no requests)
        await limits_handler.check_limits(test_user, route)
        print("Initial monthly check passed (no requests)")

        # Log 20 requests (hits the monthly limit)
        for i in range(21):
            await limits_handler.log_request(user_id, route)

        # Get current month's count
        now = datetime.now(timezone.utc)
        first_of_month = now.replace(
            day=1, hour=0, minute=0, second=0, microsecond=0
        )
        monthly_count = await limits_handler._count_requests(
            user_id, None, first_of_month
        )
        print(f"Monthly count after 20 requests: {monthly_count}")

        # This should fail as we've hit monthly_limit=20
        with pytest.raises(ValueError, match="Monthly rate limit exceeded"):
            await limits_handler.check_limits(test_user, route)

    finally:
        limits_handler.config.limits = old_limits


@pytest.mark.asyncio
async def test_user_level_override(limits_handler):
    """
    Test user-specific override limits with debug logging
    """
    user_id = UUID("47e53676-b478-5b3f-a409-234ca2164de5")
    route = "/test-route"

    # Clear existing logs first
    clear_query = f"DELETE FROM {limits_handler._get_table_name(PostgresLimitsHandler.TABLE_NAME)}"
    await limits_handler.connection_manager.execute_query(clear_query)

    test_user = User(
        id=user_id,
        email="override@example.com",
        is_active=True,
        is_verified=True,
        is_superuser=False,
        limits_overrides={
            "global_per_min": 2,
            "route_per_min": 1,
            "route_overrides": {"/test-route": {"route_per_min": 1}},
        },
    )

    # Set default limits that should be overridden
    old_limits = limits_handler.config.limits
    limits_handler.config.limits = LimitSettings(
        global_per_min=10, monthly_limit=20
    )

    # Debug: Print current limits
    print(f"\nDefault limits: {limits_handler.config.limits}")
    print(f"User overrides: {test_user.limits_overrides}")

    try:
        # First check limits (should pass as no requests yet)
        await limits_handler.check_limits(test_user, route)
        print("Initial check passed (no requests yet)")

        # Log first request
        await limits_handler.log_request(user_id, route)

        # Debug: Get current counts
        now = datetime.now(timezone.utc)
        one_min_ago = now - timedelta(minutes=1)
        global_count = await limits_handler._count_requests(
            user_id, None, one_min_ago
        )
        route_count = await limits_handler._count_requests(
            user_id, route, one_min_ago
        )
        print(f"\nAfter first request:")
        print(f"Global count: {global_count}")
        print(f"Route count: {route_count}")

        # Log second request
        await limits_handler.log_request(user_id, route)

        # This check should fail as we've hit route_per_min=1
        with pytest.raises(
            ValueError, match="Per-route per-minute rate limit exceeded"
        ):
            await limits_handler.check_limits(test_user, route)

    finally:
        # Cleanup
        limits_handler.config.limits = old_limits