r2rpy/py/tests/unit/database/test_limits.py

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

import pytest

from core.base import LimitSettings
from core.providers.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("\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


@pytest.mark.asyncio
async def test_determine_effective_limits(limits_handler):
    """Test that user-level overrides > route-level overrides > global
    defaults.

    This is a pure logic test of the 'determine_effective_limits' method.
    """
    # Setup global/base defaults
    old_limits = limits_handler.config.limits
    limits_handler.config.limits = LimitSettings(global_per_min=10,
                                                 route_per_min=5,
                                                 monthly_limit=50)

    # Setup route-level override
    route = "/some-route"
    old_route_limits = limits_handler.config.route_limits
    limits_handler.config.route_limits = {
        route: LimitSettings(global_per_min=8,
                             route_per_min=3,
                             monthly_limit=30)
    }

    # Setup user-level override
    test_user = User(
        id=uuid.uuid4(),
        email="test@example.com",
        is_active=True,
        is_verified=True,
        is_superuser=False,
        limits_overrides={
            "global_per_min": 6,  # should override
            "route_overrides": {
                route: {
                    "route_per_min": 2
                }  # should override
            },
        },
    )

    try:
        effective = limits_handler.determine_effective_limits(test_user, route)

        # Check final / effective limits
        # Global limit overridden to 6
        assert effective.global_per_min == 6, (
            "User-level global override not applied")

        # route_per_min should be overridden to 2 (not the route-level 3)
        assert effective.route_per_min == 2, (
            "User-level route override not applied")

        # monthly_limit from route-level override is 30, user didn't override it, so it should stay 30
        assert effective.monthly_limit == 30, (
            "Route-level monthly override not applied")
    finally:
        # revert changes
        limits_handler.config.limits = old_limits
        limits_handler.config.route_limits = old_route_limits


@pytest.mark.asyncio
async def test_separate_route_usage_is_isolated(limits_handler):
    """Confirm that calls to /routeA do NOT increment the per-route usage for
    /routeB, and vice-versa."""
    # 1) Clear existing logs
    clear_query = f"DELETE FROM {limits_handler._get_table_name(limits_handler.TABLE_NAME)}"
    await limits_handler.connection_manager.execute_query(clear_query)

    # 2) Setup user & routes
    import uuid

    from shared.abstractions import User

    user_id = uuid.uuid4()
    routeA = "/v3/retrieval/rag"
    routeB = "/v3/retrieval/search"

    test_user = User(
        id=user_id,
        email="test@example.com",
        is_active=True,
        is_verified=True,
        is_superuser=False,
        limits_overrides=None,
    )

    # 3) Insert some logs for routeA only
    for _ in range(3):
        await limits_handler.log_request(user_id, routeA)

    # 4) Check usage for routeA → Should be 3 in last minute
    now = datetime.now(timezone.utc)
    one_min_ago = now - timedelta(minutes=1)
    routeA_count = await limits_handler._count_requests(
        user_id, routeA, one_min_ago)
    assert routeA_count == 3, f"Expected 3 for routeA, got {routeA_count}"

    # 5) Check usage for routeB → Should be 0
    routeB_count = await limits_handler._count_requests(
        user_id, routeB, one_min_ago)
    assert routeB_count == 0, f"Expected 0 for routeB, got {routeB_count}"

    # 6) Insert some logs for routeB only
    for _ in range(2):
        await limits_handler.log_request(user_id, routeB)

    # 7) Recheck usage
    routeA_count_after = await limits_handler._count_requests(
        user_id, routeA, one_min_ago)
    routeB_count_after = await limits_handler._count_requests(
        user_id, routeB, one_min_ago)
    assert routeA_count_after == 3, (
        f"RouteA usage changed unexpectedly: {routeA_count_after}")
    assert routeB_count_after == 2, (
        f"RouteB usage is wrong: {routeB_count_after}")


# @pytest.mark.asyncio
# async def test_check_limits_multiple_routes(limits_handler):
#     """
#     Demonstrates that routeA calls do not count against routeB's per-minute limit.
#     """
#     # Clear logs
#     clear_query = f"DELETE FROM {limits_handler._get_table_name(limits_handler.TABLE_NAME)}"
#     await limits_handler.connection_manager.execute_query(clear_query)

#     import uuid
#     from shared.abstractions import User
#     user_id = uuid.uuid4()
#     routeA = "/v3/retrieval/rag"
#     routeB = "/v3/retrieval/search"

#     # Suppose routeA has a limit of 2/min, routeB has a limit of 3/min
#     # (You can do this by setting config.route_limits[routeA].route_per_min, etc.)
#     # Or just rely on your global config if needed.

#     test_user = User(
#         id=user_id,
#         email="test@example.com",
#         is_active=True,
#         is_verified=True,
#         is_superuser=False,
#         limits_overrides=None,
#     )

#     # 1) Make 2 calls to routeA
#     await limits_handler.check_limits(test_user, routeA)
#     await limits_handler.log_request(user_id, routeA)

#     await limits_handler.check_limits(test_user, routeA)
#     await limits_handler.log_request(user_id, routeA)
#     await limits_handler.check_limits(test_user, routeA)
#     await limits_handler.log_request(user_id, routeA)

#     # 2) Confirm next call to routeA fails if the limit is 2/min
#     with pytest.raises(ValueError, match="Per-route per-minute rate limit exceeded"):
#         await limits_handler.check_limits(test_user, routeA)

#     # 3) Meanwhile, routeB usage should be unaffected
#     #    We can still do 3 calls to routeB (assuming route_per_min=3).
#     await limits_handler.check_limits(test_user, routeB)
#     await limits_handler.log_request(user_id, routeB)
#     await limits_handler.check_limits(test_user, routeB)
#     await limits_handler.log_request(user_id, routeB)
#     await limits_handler.check_limits(test_user, routeB)
#     await limits_handler.log_request(user_id, routeB)


@pytest.mark.asyncio
async def test_route_specific_monthly_usage(limits_handler):
    """Confirm that monthly usage is tracked per-route and doesn't get
    incremented by calls to other routes."""
    # 1) Clear existing logs
    clear_query = f"DELETE FROM {limits_handler._get_table_name(limits_handler.TABLE_NAME)}"
    await limits_handler.connection_manager.execute_query(clear_query)

    # 2) Setup
    user_id = uuid.uuid4()
    routeA = "/v3/retrieval/rag"
    routeB = "/v3/retrieval/search"
    test_user = User(
        id=user_id,
        email="test_monthly_routes@example.com",
        is_active=True,
        is_verified=True,
        is_superuser=False,
        limits_overrides=None,
    )

    # 3) Log 5 requests for routeA
    for _ in range(5):
        await limits_handler.log_request(user_id, routeA)

    # 4) Check monthly usage for routeA => should be 5
    routeA_monthly = await limits_handler._count_monthly_requests(
        user_id, routeA)
    assert routeA_monthly == 5, f"Expected 5 for routeA, got {routeA_monthly}"

    # routeB => should still be 0
    routeB_monthly = await limits_handler._count_monthly_requests(
        user_id, routeB)
    assert routeB_monthly == 0, f"Expected 0 for routeB, got {routeB_monthly}"

    # 5) Now log 3 requests for routeB
    for _ in range(3):
        await limits_handler.log_request(user_id, routeB)

    # Re-check usage
    routeA_monthly_after = await limits_handler._count_monthly_requests(
        user_id, routeA)
    routeB_monthly_after = await limits_handler._count_monthly_requests(
        user_id, routeB)
    assert routeA_monthly_after == 5, (
        f"RouteA usage changed unexpectedly: {routeA_monthly_after}")
    assert routeB_monthly_after == 3, (
        f"RouteB usage is wrong: {routeB_monthly_after}")

    # Additionally confirm total usage across all routes
    global_monthly = await limits_handler._count_monthly_requests(user_id,
                                                                  route=None)
    assert global_monthly == 8, (
        f"Expected total of 8 monthly requests, got {global_monthly}")