replypier/backend/tests/utils/test_rate_limit.py

"""
Copyright 2024-2026 ReplyPier

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import pytest
from unittest.mock import MagicMock, patch
from fastapi import HTTPException, Request
import sys
from pathlib import Path

# Add the backend directory to Python path
backend_dir = str(Path(__file__).parent.parent.parent)
if backend_dir not in sys.path:
    sys.path.append(backend_dir)

# Mock Redis and FastAPI dependencies before importing app modules
with patch('redis.Redis'), \
     patch('fastapi.staticfiles.StaticFiles'):
    from app.utils.rate_limit import (
        format_time_remaining,
        check_rate_limit,
        rate_limit,
        limit_instruction_generation,
        RateLimitConfig
    )

def test_format_time_remaining():
    """Test time formatting for different durations"""
    # Test hours and minutes
    assert format_time_remaining(3665) == "1 hours and 1 minutes"
    
    # Test minutes and seconds
    assert format_time_remaining(125) == "2 minutes and 5 seconds"
    
    # Test only seconds
    assert format_time_remaining(45) == "45 seconds"
    
    # Test only hours
    assert format_time_remaining(7200) == "2 hours"
    
    # Test only minutes
    assert format_time_remaining(180) == "3 minutes"

def test_check_rate_limit_first_request():
    """Test rate limit check for first request"""
    mock_redis = MagicMock()
    mock_redis.get.return_value = None
    
    check_rate_limit(mock_redis, "test_key", 10, 3600)
    
    mock_redis.setex.assert_called_once_with("test_key", 3600, 1)

def test_check_rate_limit_under_limit():
    """Test rate limit check when under the limit"""
    mock_redis = MagicMock()
    mock_redis.get.return_value = "5"
    
    check_rate_limit(mock_redis, "test_key", 10, 3600)
    
    mock_redis.incr.assert_called_once_with("test_key")

def test_check_rate_limit_exceeded():
    """Test rate limit check when limit is exceeded"""
    mock_redis = MagicMock()
    mock_redis.get.return_value = "10"
    mock_redis.ttl.return_value = 1800  # 30 minutes remaining
    
    with pytest.raises(HTTPException) as exc_info:
        check_rate_limit(mock_redis, "test_key", 10, 3600)
    
    assert exc_info.value.status_code == 429
    assert "30 minutes" in exc_info.value.detail

def test_check_rate_limit_no_redis():
    """Test rate limit check when Redis is not available"""
    result = check_rate_limit(None, "test_key", 10, 3600)
    assert result is None

@pytest.mark.skip(reason="Need to modify source code to handle dict objects properly")
@pytest.mark.asyncio
async def test_rate_limit_decorator():
    """Test rate limit decorator functionality"""
    mock_redis = MagicMock()
    mock_redis.get.return_value = None

    # Mock the request with proper client IP
    mock_request = MagicMock(spec=Request)
    mock_request.client = MagicMock()
    mock_request.client.host = "127.0.0.1"

    # Create a current_user dictionary with organization_id
    mock_user = {"organization_id": "org123"}

    # Create a decorated test function
    @rate_limit(limit=10, window=3600, key_prefix="test")
    async def test_func(request, current_user=None):
        return "success"

    with patch('app.utils.rate_limit.get_redis', return_value=mock_redis), \
         patch('app.core.config.settings.REDIS_ENABLED', True):
        result = await test_func(mock_request, current_user=mock_user)
        assert result == "success"

@pytest.mark.skip(reason="Need to modify source code to handle dict objects properly")
@pytest.mark.asyncio
async def test_rate_limit_decorator_exceeded():
    """Test rate limit decorator when limit is exceeded"""
    mock_redis = MagicMock()
    mock_redis.get.return_value = "10"  # At limit
    mock_redis.ttl.return_value = 1800  # 30 minutes remaining

    # Mock the request with proper client IP
    mock_request = MagicMock(spec=Request)
    mock_request.client = MagicMock()
    mock_request.client.host = "127.0.0.1"

    # Create a current_user dictionary with organization_id
    mock_user = {"organization_id": "org123"}

    # Create a decorated test function
    @rate_limit(limit=10, window=3600, key_prefix="test")
    async def test_func(request, current_user=None):
        return "success"

    with patch('app.utils.rate_limit.get_redis', return_value=mock_redis), \
         patch('app.core.config.settings.REDIS_ENABLED', True):
        with pytest.raises(HTTPException) as exc_info:
            await test_func(mock_request, current_user=mock_user)
        assert exc_info.value.status_code == 429

@pytest.mark.skip(reason="Need to modify source code to handle dict objects properly")
@pytest.mark.asyncio
async def test_limit_instruction_generation():
    """Test instruction generation rate limit decorator"""
    mock_redis = MagicMock()
    mock_redis.get.return_value = None

    # Mock the request with proper client IP
    mock_request = MagicMock(spec=Request)
    mock_request.client = MagicMock()
    mock_request.client.host = "127.0.0.1"

    # Create a current_user dictionary with organization_id
    mock_user = {"organization_id": "org123"}

    # Create a decorated test function
    @limit_instruction_generation
    async def test_func(request, current_user=None):
        return "success"

    with patch('app.utils.rate_limit.get_redis', return_value=mock_redis), \
         patch('app.core.config.settings.REDIS_ENABLED', True):
        result = await test_func(mock_request, current_user=mock_user)
        assert result == "success"
        
        # Verify Redis interactions with correct limit values
        mock_redis.setex.assert_called_once_with(
            f"{RateLimitConfig.INSTRUCTION_GEN_PREFIX}:org123:unknown",
            RateLimitConfig.INSTRUCTION_GEN_WINDOW,
            1
        )

@pytest.mark.skip(reason="Need to modify source code to handle dict objects properly")
@pytest.mark.asyncio
async def test_rate_limit_with_org_id():
    """Test rate limit decorator with organization ID"""
    mock_redis = MagicMock()
    mock_redis.get.return_value = None
    
    # Mock the request with proper client IP
    mock_request = MagicMock(spec=Request)
    mock_request.client = MagicMock()
    mock_request.client.host = "127.0.0.1"
    
    # Create a decorated test function
    @rate_limit(limit=10, window=3600, key_prefix="test")
    async def test_func(request, current_user=None):
        return "success"
    
    with patch('app.utils.rate_limit.get_redis', return_value=mock_redis), \
         patch('app.core.config.settings.REDIS_ENABLED', True):
        # Create a current_user dictionary with organization_id
        current_user = {"organization_id": "org123"}
        
        result = await test_func(mock_request, current_user=current_user)
        assert result == "success"
        
        # Verify Redis key includes org ID but with "unknown" for client IP since that's how the implementation works
        mock_redis.setex.assert_called_with(
            "test:org123:unknown",
            3600,
            1
        )