learningcircuit--local-deep-research/tests/rate_limiting/test_llm_wrapper_extended.py

"""
Extended tests for rate_limiting/llm/wrapper.py — active rate limiting paths.

The existing test_llm_rate_limiting.py only covers disabled rate limiting.
Two tests are skipped with "Rate limiting is disabled by default".
These tests exercise the actual rate-limiting code paths.

Tests cover:
- AdaptiveLLMWait.__call__: tracker wait, retry-after merging, error recording
- RateLimitedLLMWrapper._check_if_local_model: provider + URL detection
- RateLimitedLLMWrapper._do_invoke: error wrapping
- RateLimitedLLMWrapper.invoke with rate limiting active
- __str__ and __repr__
"""

from unittest.mock import Mock

import pytest


WRAPPER_MODULE = (
    "local_deep_research.web_search_engines.rate_limiting.llm.wrapper"
)


# ---------------------------------------------------------------------------
# AdaptiveLLMWait.__call__
# ---------------------------------------------------------------------------


class TestAdaptiveLLMWait:
    """Tests for AdaptiveLLMWait wait strategy."""

    def test_returns_tracker_wait_time_no_error(self):
        """Returns tracker wait time when no error."""
        from local_deep_research.web_search_engines.rate_limiting.llm.wrapper import (
            AdaptiveLLMWait,
        )

        mock_tracker = Mock()
        mock_tracker.get_wait_time.return_value = 5.0

        wait = AdaptiveLLMWait(mock_tracker, "openai-api")

        # Simulate retry_state with no failure
        retry_state = Mock()
        retry_state.outcome = None

        result = wait(retry_state)

        assert result == 5.0
        mock_tracker.get_wait_time.assert_called_once_with("openai-api")

    def test_returns_max_of_tracker_and_retry_after(self):
        """Returns max(tracker_wait, retry_after) when error has retry-after."""
        from local_deep_research.web_search_engines.rate_limiting.llm.wrapper import (
            AdaptiveLLMWait,
        )

        mock_tracker = Mock()
        mock_tracker.get_wait_time.return_value = 2.0

        wait = AdaptiveLLMWait(mock_tracker, "openai-api")

        # First call: record an error with retry-after
        error = Mock()
        error.response = Mock()
        error.response.headers = {"Retry-After": "10"}
        error.__str__ = lambda self: "rate limit"

        retry_state = Mock()
        retry_state.outcome = Mock()
        retry_state.outcome.failed = True
        retry_state.outcome.exception.return_value = error

        result = wait(retry_state)

        # Should be max(2.0, 10.0) = 10.0
        assert result == 10.0

    def test_records_last_error_from_failed_outcome(self):
        """Records last_error from failed outcome."""
        from local_deep_research.web_search_engines.rate_limiting.llm.wrapper import (
            AdaptiveLLMWait,
        )

        mock_tracker = Mock()
        mock_tracker.get_wait_time.return_value = 1.0

        wait = AdaptiveLLMWait(mock_tracker, "test")

        error = Exception("rate limit hit")

        retry_state = Mock()
        retry_state.outcome = Mock()
        retry_state.outcome.failed = True
        retry_state.outcome.exception.return_value = error

        wait(retry_state)

        assert wait.last_error is error

    def test_uses_last_error_across_calls(self):
        """last_error persists and is used in subsequent calls."""
        from local_deep_research.web_search_engines.rate_limiting.llm.wrapper import (
            AdaptiveLLMWait,
        )

        mock_tracker = Mock()
        mock_tracker.get_wait_time.return_value = 1.0

        wait = AdaptiveLLMWait(mock_tracker, "test")

        # First call: record error with retry-after in message
        error = Exception("retry after 30 seconds")
        retry_state = Mock()
        retry_state.outcome = Mock()
        retry_state.outcome.failed = True
        retry_state.outcome.exception.return_value = error
        result1 = wait(retry_state)
        assert result1 == 30.0

        # Second call: no new error, but last_error still has retry-after
        retry_state2 = Mock()
        retry_state2.outcome = None
        result2 = wait(retry_state2)
        assert result2 == 30.0


# ---------------------------------------------------------------------------
# RateLimitedLLMWrapper._check_if_local_model
# ---------------------------------------------------------------------------


class TestCheckIfLocalModel:
    """Tests for _check_if_local_model."""

    def test_true_for_local_providers(self):
        """Returns True for known local provider names."""
        from local_deep_research.web_search_engines.rate_limiting.llm.wrapper import (
            create_rate_limited_llm_wrapper,
        )

        for provider in [
            "ollama",
            "lmstudio",
            "llamacpp",
            "local",
            "none",
        ]:
            base_llm = Mock()
            wrapper = create_rate_limited_llm_wrapper(
                base_llm, provider=provider
            )
            assert wrapper._check_if_local_model() is True, (
                f"Failed for {provider}"
            )

    def test_true_for_localhost_base_url(self):
        """Returns True when base_url contains localhost."""
        from local_deep_research.web_search_engines.rate_limiting.llm.wrapper import (
            create_rate_limited_llm_wrapper,
        )

        base_llm = Mock()
        base_llm.base_url = "http://localhost:11434/v1"
        wrapper = create_rate_limited_llm_wrapper(base_llm, provider="custom")
        assert wrapper._check_if_local_model() is True

    def test_true_for_127_0_0_1_base_url(self):
        """Returns True when base_url contains 127.0.0.1."""
        from local_deep_research.web_search_engines.rate_limiting.llm.wrapper import (
            create_rate_limited_llm_wrapper,
        )

        base_llm = Mock()
        base_llm.base_url = "http://127.0.0.1:8080/api"
        wrapper = create_rate_limited_llm_wrapper(base_llm, provider="custom")
        assert wrapper._check_if_local_model() is True

    def test_false_for_remote_urls(self):
        """Returns False for remote URLs and providers."""
        from local_deep_research.web_search_engines.rate_limiting.llm.wrapper import (
            create_rate_limited_llm_wrapper,
        )

        base_llm = Mock()
        base_llm.base_url = "https://api.openai.com/v1"
        wrapper = create_rate_limited_llm_wrapper(base_llm, provider="openai")
        assert wrapper._check_if_local_model() is False


# ---------------------------------------------------------------------------
# RateLimitedLLMWrapper._do_invoke
# ---------------------------------------------------------------------------


class TestDoInvoke:
    """Tests for _do_invoke."""

    def test_rate_limit_error_wrapped_into_rate_limit_error(self):
        """Rate limit errors are wrapped into RateLimitError."""
        from local_deep_research.web_search_engines.rate_limiting.llm.wrapper import (
            create_rate_limited_llm_wrapper,
        )
        from local_deep_research.web_search_engines.rate_limiting.exceptions import (
            RateLimitError,
        )

        base_llm = Mock()
        base_llm.invoke.side_effect = Exception(
            "429 Too Many Requests rate limit"
        )
        wrapper = create_rate_limited_llm_wrapper(base_llm)

        with pytest.raises(RateLimitError, match="LLM rate limit"):
            wrapper._do_invoke("test prompt")

    def test_non_rate_limit_error_passed_through(self):
        """Non-rate-limit errors pass through unchanged."""
        from local_deep_research.web_search_engines.rate_limiting.llm.wrapper import (
            create_rate_limited_llm_wrapper,
        )

        base_llm = Mock()
        base_llm.invoke.side_effect = ValueError("Invalid input format")
        wrapper = create_rate_limited_llm_wrapper(base_llm)

        with pytest.raises(ValueError, match="Invalid input format"):
            wrapper._do_invoke("test prompt")

    def test_successful_invoke_returns_result(self):
        """Successful invoke returns result."""
        from local_deep_research.web_search_engines.rate_limiting.llm.wrapper import (
            create_rate_limited_llm_wrapper,
        )

        base_llm = Mock()
        base_llm.invoke.return_value = "LLM response"
        wrapper = create_rate_limited_llm_wrapper(base_llm)

        result = wrapper._do_invoke("test prompt")
        assert result == "LLM response"


# ---------------------------------------------------------------------------
# RateLimitedLLMWrapper.invoke with rate limiting active
# ---------------------------------------------------------------------------


class TestInvokeWithRateLimiting:
    """Tests for invoke() with rate limiting enabled."""

    def test_successful_first_attempt_records_success(self):
        """Successful first attempt records success in tracker."""
        from local_deep_research.web_search_engines.rate_limiting.llm.wrapper import (
            create_rate_limited_llm_wrapper,
        )

        base_llm = Mock()
        base_llm.invoke.return_value = "response"

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

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

        # Force-enable rate limiting
        wrapper.rate_limiter = mock_tracker

        result = wrapper.invoke("prompt")

        assert result == "response"
        mock_tracker.record_outcome.assert_called_once_with(
            engine_type=wrapper._get_rate_limit_key(),
            wait_time=0,
            success=True,
            retry_count=0,
        )

    def test_invoke_without_rate_limiter_delegates_directly(self):
        """Without rate limiter, invoke delegates to _do_invoke."""
        from local_deep_research.web_search_engines.rate_limiting.llm.wrapper import (
            create_rate_limited_llm_wrapper,
        )

        base_llm = Mock()
        base_llm.invoke.return_value = "direct response"

        wrapper = create_rate_limited_llm_wrapper(base_llm)
        assert wrapper.rate_limiter is None

        result = wrapper.invoke("prompt")
        assert result == "direct response"
        base_llm.invoke.assert_called_once_with("prompt")

    def test_rate_limit_error_records_failure(self):
        """Rate limit error is recorded as failure in tracker."""
        from local_deep_research.web_search_engines.rate_limiting.llm.wrapper import (
            create_rate_limited_llm_wrapper,
        )

        base_llm = Mock()
        base_llm.invoke.side_effect = Exception("429 rate limit exceeded")

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

        wrapper = create_rate_limited_llm_wrapper(base_llm, provider="openai")
        wrapper.rate_limiter = mock_tracker

        # tenacity wraps the final error in RetryError after exhausting retries
        with pytest.raises(Exception):
            wrapper.invoke("prompt")

        # Should have recorded a failure
        failure_calls = [
            c
            for c in mock_tracker.record_outcome.call_args_list
            if c.kwargs.get("success") is False
        ]
        assert len(failure_calls) >= 1


# ---------------------------------------------------------------------------
# __str__ and __repr__
# ---------------------------------------------------------------------------


class TestStringRepresentations:
    """Tests for __str__ and __repr__."""

    def test_str_format(self):
        """__str__ returns expected format."""
        from local_deep_research.web_search_engines.rate_limiting.llm.wrapper import (
            create_rate_limited_llm_wrapper,
        )

        base_llm = Mock()
        base_llm.__str__ = lambda self: "MockLLM"
        wrapper = create_rate_limited_llm_wrapper(base_llm)

        assert str(wrapper) == "RateLimited(MockLLM)"

    def test_repr_format(self):
        """__repr__ returns expected format."""
        from local_deep_research.web_search_engines.rate_limiting.llm.wrapper import (
            create_rate_limited_llm_wrapper,
        )

        base_llm = Mock()
        wrapper = create_rate_limited_llm_wrapper(base_llm)

        result = repr(wrapper)
        assert result.startswith("RateLimitedLLMWrapper(")