learningcircuit--local-deep-research/tests/test_llm_rate_limiting.py

"""
Tests for LLM rate limiting functionality.
"""

import pytest
from unittest.mock import Mock, patch
from langchain_core.messages import AIMessage

from tests.test_utils import add_src_to_path

# Add src to path
add_src_to_path()

from local_deep_research.web_search_engines.rate_limiting.llm import (  # noqa: E402
    create_rate_limited_llm_wrapper,
    is_llm_rate_limit_error,
)


class TestLLMRateLimitDetection:
    """Test rate limit error detection for various LLM providers."""

    def test_detects_429_status_code(self):
        """Test detection of HTTP 429 status code."""
        error = Mock()
        error.response = Mock()
        error.response.status_code = 429

        assert is_llm_rate_limit_error(error) is True

    def test_detects_rate_limit_messages(self):
        """Test detection of common rate limit error messages."""
        test_messages = [
            "Error: 429 Resource has been exhausted (e.g. check quota).",
            "Rate limit exceeded. Please try again later.",
            "Too many requests. Quota exceeded.",
            "API rate limit: maximum requests per minute exceeded",
            "Threshold for requests has been reached",
        ]

        for message in test_messages:
            error = Exception(message)
            assert is_llm_rate_limit_error(error) is True, (
                f"Failed to detect: {message}"
            )

    def test_does_not_detect_non_rate_limit_errors(self):
        """Test that non-rate limit errors are not detected."""
        test_messages = [
            "Model not found",
            "Invalid API key",
            "Connection refused",
            "Internal server error",
        ]

        for message in test_messages:
            error = Exception(message)
            assert is_llm_rate_limit_error(error) is False, (
                f"Incorrectly detected: {message}"
            )


class TestRateLimitedLLMWrapper:
    """Test the rate limited LLM wrapper functionality."""

    @pytest.fixture
    def mock_llm(self):
        """Create a mock LLM for testing."""
        llm = Mock()
        llm.model_name = "test-model"
        llm.base_url = "https://api.example.com"
        llm.invoke = Mock(return_value=AIMessage(content="Test response"))
        return llm

    @pytest.fixture
    def mock_db_settings(self):
        """Mock database settings."""
        # Rate limiting is now disabled by default in the wrapper
        # This fixture is no longer needed but kept for compatibility
        yield None

    def test_wrapper_creation_with_rate_limiting_disabled_by_default(
        self, mock_llm, mock_db_settings
    ):
        """Test wrapper creation - rate limiting is now disabled by default."""
        wrapper = create_rate_limited_llm_wrapper(mock_llm, provider="openai")

        assert wrapper is not None
        assert hasattr(wrapper, "invoke")
        assert wrapper.base_llm == mock_llm
        assert wrapper.provider == "openai"
        # Rate limiting is disabled by default now
        assert wrapper.rate_limiter is None

    def test_wrapper_creation_always_disabled(self, mock_llm):
        """Test wrapper creation - rate limiting is always disabled now."""
        # No need to patch, rate limiting is disabled by default
        wrapper = create_rate_limited_llm_wrapper(mock_llm, provider="openai")

        assert wrapper.rate_limiter is None

    def test_local_providers_skip_rate_limiting(
        self, mock_llm, mock_db_settings
    ):
        """Test that local providers skip rate limiting even when enabled."""
        local_providers = ["ollama", "lmstudio", "llamacpp"]

        for provider in local_providers:
            wrapper = create_rate_limited_llm_wrapper(
                mock_llm, provider=provider
            )
            assert wrapper.rate_limiter is None, (
                f"Rate limiting should be skipped for {provider}"
            )

    def test_rate_limit_key_generation(self, mock_llm, mock_db_settings):
        """Test the generation of rate limit keys."""
        wrapper = create_rate_limited_llm_wrapper(mock_llm, provider="openai")

        key = wrapper._get_rate_limit_key()
        assert key == "openai-api.example.com-test-model"

    def test_invoke_without_rate_limiting(self, mock_llm):
        """Test invoke when rate limiting is disabled (always the case now)."""
        # No need to patch, rate limiting is disabled by default
        wrapper = create_rate_limited_llm_wrapper(mock_llm, provider="openai")
        response = wrapper.invoke("Test prompt")

        assert response.content == "Test response"
        mock_llm.invoke.assert_called_once_with("Test prompt")

    def test_attribute_passthrough(self, mock_llm, mock_db_settings):
        """Test that attributes are passed through to the base LLM."""
        mock_llm.custom_attribute = "test_value"

        wrapper = create_rate_limited_llm_wrapper(mock_llm, provider="openai")

        assert wrapper.custom_attribute == "test_value"
        assert wrapper.model_name == "test-model"


class TestIntegrationWithTracker:
    """Test integration with the adaptive rate limit tracker."""

    @patch(
        "local_deep_research.web_search_engines.rate_limiting.llm.wrapper.get_tracker"
    )
    def test_tracker_not_used_when_disabled(self, mock_get_tracker):
        """Test that tracker is not used since rate limiting is disabled."""
        # Rate limiting is disabled by default, so tracker shouldn't be called

        mock_tracker = Mock()
        mock_tracker.get_wait_time.return_value = 0
        mock_get_tracker.return_value = mock_tracker

        mock_llm = Mock()
        mock_llm.model_name = "test-model"
        mock_llm.base_url = "https://api.example.com"
        mock_llm.invoke.return_value = AIMessage(content="Success")

        wrapper = create_rate_limited_llm_wrapper(mock_llm, provider="openai")

        # Verify rate limiter is None (disabled)
        assert wrapper.rate_limiter is None

        # get_tracker should not be called since rate limiting is disabled
        mock_get_tracker.assert_not_called()

        # But invoke should still work
        result = wrapper.invoke("Test prompt")
        assert result.content == "Success"


if __name__ == "__main__":
    pytest.main([__file__, "-v"])