""" Unit tests for rate-limit retry behaviour across all LLM call sites. Covers acceptance criteria: - 5-retry exhaustion raises RateLimitError - Retry-After header drives the sleep duration - Partial retries then success returns the correct result - ConnectError / TimeoutException in OllamaEmbedder does NOT trigger the 429 retry path (the two concerns are composed without double-retrying) """ from __future__ import annotations import json import sys import types from typing import Any from unittest.mock import AsyncMock, MagicMock, patch from uuid import uuid4 import anthropic as anthropic_mod import httpx import pytest import pytest_asyncio # noqa: F401 - registers asyncio mode # --------------------------------------------------------------------------- # Ensure piragi stubs are present before the optimal_brain modules are imported # --------------------------------------------------------------------------- _PIRAGI_STUB_NAMES = ( "piragi", "piragi.types", "piragi.stores", "piragi.stores.postgres", "piragi.chunking", "piragi.semantic_chunking", ) def _stub_piragi() -> None: mock = MagicMock() for name in _PIRAGI_STUB_NAMES: if name not in sys.modules: mod = types.ModuleType(name) mod.__dict__.update( { "AsyncRagi": mock, "Citation": mock, "Document": mock, "Chunk": mock, "PostgresStore": mock, } ) sys.modules[name] = mod _stub_piragi() # --------------------------------------------------------------------------- # Module imports (after stubs are injected) # --------------------------------------------------------------------------- from roboco.models.extraction import ExtractionContext # noqa: E402 from roboco.models.optimal import IndexType # noqa: E402 from roboco.services.exceptions import ( # noqa: E402 MAX_RATE_LIMIT_RETRIES, RateLimitError, parse_retry_after_header, ) from roboco.services.extraction import ExtractionService # noqa: E402 from roboco.services.optimal_brain.indexes.journals import ( # noqa: E402 JournalsIndexPlugin, ) from roboco.services.optimal_brain.mentor import MentorService # noqa: E402 from roboco.services.optimal_brain.ollama_embedder import ( # noqa: E402 MAX_RETRIES, OllamaConnectionError, OllamaEmbedder, ) from roboco.services.optimal_brain.validator import ValidatorService # noqa: E402 # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- _RETRY_AFTER_FLOAT = 30.0 _RETRY_AFTER_FLOAT_2 = 2.5 _RETRY_AFTER_7 = 7.0 _RETRY_AFTER_9 = 9.0 _RETRY_AFTER_12 = 12.0 _RETRY_AFTER_5 = 5.0 _EMBED_DIM = 4 # zero-vector dimension in test responses _CALLS_2RL_1_SUCCESS = 3 # 2 rate-limit errors then 1 success _EMBED_PATH = ( "roboco.services.optimal_brain.ollama_embedder.OllamaEmbedder._create_async_client" ) # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _make_response( status_code: int, body: Any = None, retry_after: str | None = None, ) -> httpx.Response: """Build a minimal httpx.Response for use in mocks.""" headers: dict[str, str] = {} if retry_after is not None: headers["retry-after"] = retry_after content = json.dumps(body or {}).encode() return httpx.Response( status_code=status_code, headers=headers, content=content, ) def _success_embed_response(n: int = 1) -> httpx.Response: """Return a valid Ollama /api/embed response with *n* zero-vectors.""" body = {"embeddings": [[0.0] * _EMBED_DIM] * n} return _make_response(200, body) def _429_response(retry_after: str | None = None) -> httpx.Response: return _make_response(429, {"error": "rate limited"}, retry_after=retry_after) def _make_async_client_mock(post_return: Any = None) -> AsyncMock: """Return an async context manager mock exposing `.post`.""" client = AsyncMock() client.__aenter__ = AsyncMock(return_value=client) client.__aexit__ = AsyncMock(return_value=False) if post_return is not None: client.post = AsyncMock(return_value=post_return) return client def _make_extraction_context() -> ExtractionContext: return ExtractionContext( content="This is a test message that is long enough.", agent_id=uuid4(), channel_id=uuid4(), session_id=uuid4(), group_id=uuid4(), ) def _make_anthropic_rl_exc( retry_after: str | None = None, ) -> anthropic_mod.RateLimitError: """Build a minimal anthropic.RateLimitError.""" fake_resp = MagicMock() fake_resp.headers = {} if retry_after is None else {"retry-after": retry_after} return anthropic_mod.RateLimitError( message="rate limit", response=fake_resp, body={}, ) def _make_journal_plugin() -> JournalsIndexPlugin: """Create a minimal JournalsIndexPlugin without initialising piragi.""" plugin = JournalsIndexPlugin.__new__(JournalsIndexPlugin) plugin._config = MagicMock() plugin._config.llm_base_url = "http://ollama-test:11434/v1" plugin._config.llm_model = "glm-5:cloud" plugin._ragi = MagicMock() plugin._initialized = True return plugin def _make_search_outcome(content: str = "context text") -> Any: mock_outcome = MagicMock() mock_outcome.success = True mock_outcome.results = [ MagicMock(content=content, source="src", score=0.9, index_type=None) ] return mock_outcome def _make_sources() -> list[Any]: return [ MagicMock(content="ctx", source="s", score=0.9, index_type=IndexType.JOURNALS) ] def _make_standards() -> list[Any]: s = MagicMock() s.content = "### PY-001: Use Type Hints\nMust add return type annotations." return [s] # =========================================================================== # 1. parse_retry_after_header # =========================================================================== class TestParseRetryAfterHeader: def test_integer_seconds(self) -> None: resp = _429_response(retry_after="30") assert parse_retry_after_header(resp) == _RETRY_AFTER_FLOAT def test_float_seconds(self) -> None: resp = _429_response(retry_after="2.5") assert parse_retry_after_header(resp) == _RETRY_AFTER_FLOAT_2 def test_missing_header_returns_none(self) -> None: resp = _make_response(429) assert parse_retry_after_header(resp) is None def test_non_numeric_returns_none(self) -> None: resp = _429_response(retry_after="Wed, 21 Oct 2015 07:28:00 GMT") assert parse_retry_after_header(resp) is None # =========================================================================== # 2. RateLimitError class # =========================================================================== class TestRateLimitError: def test_fields(self) -> None: err = RateLimitError(provider="anthropic", retry_after=_RETRY_AFTER_FLOAT) assert err.provider == "anthropic" assert err.retry_after == _RETRY_AFTER_FLOAT def test_message_includes_provider(self) -> None: err = RateLimitError(provider="ollama") assert "ollama" in str(err) def test_message_includes_retry_after_when_set(self) -> None: err = RateLimitError(provider="ollama", retry_after=15.0) assert "15" in str(err) def test_none_retry_after(self) -> None: err = RateLimitError(provider="anthropic", retry_after=None) assert err.retry_after is None # =========================================================================== # 3. OllamaEmbedder - async path (aembed_query) # =========================================================================== class TestOllamaEmbedderAembed: """Tests for the async aembed_query method.""" def _make_embedder(self) -> OllamaEmbedder: return OllamaEmbedder(base_url="http://ollama-test:11434") async def test_five_consecutive_429s_raise_rate_limit_error(self) -> None: """After MAX_RATE_LIMIT_RETRIES attempts all 429 -> RateLimitError.""" embedder = self._make_embedder() mock_c = _make_async_client_mock(post_return=_429_response()) with ( patch(_EMBED_PATH, return_value=mock_c), patch("asyncio.sleep", new_callable=AsyncMock), pytest.raises(RateLimitError) as exc_info, ): await embedder.aembed_query("hello") assert exc_info.value.provider == "ollama" assert mock_c.post.call_count == MAX_RATE_LIMIT_RETRIES async def test_retry_after_header_respected_as_sleep_duration(self) -> None: """Retry-After: 7 -> asyncio.sleep(7.0) on each inter-attempt gap.""" embedder = self._make_embedder() sleep_calls: list[float] = [] async def _fake_sleep(secs: float) -> None: sleep_calls.append(secs) mock_c = _make_async_client_mock(post_return=_429_response(retry_after="7")) with ( patch(_EMBED_PATH, return_value=mock_c), patch("asyncio.sleep", side_effect=_fake_sleep), pytest.raises(RateLimitError), ): await embedder.aembed_query("hello") assert all(s == _RETRY_AFTER_7 for s in sleep_calls) assert len(sleep_calls) == MAX_RATE_LIMIT_RETRIES - 1 async def test_partial_retries_then_success(self) -> None: """Two 429s then a 200 -> returns the embedding list.""" embedder = self._make_embedder() mock_c = AsyncMock() mock_c.__aenter__ = AsyncMock(return_value=mock_c) mock_c.__aexit__ = AsyncMock(return_value=False) mock_c.post = AsyncMock( side_effect=[ _429_response(), _429_response(), _success_embed_response(), ] ) with ( patch(_EMBED_PATH, return_value=mock_c), patch("asyncio.sleep", new_callable=AsyncMock), ): result = await embedder.aembed_query("hello") assert isinstance(result, list) assert len(result) == _EMBED_DIM async def test_connect_error_does_not_trigger_429_retry_path( self, ) -> None: """ConnectError -> OllamaConnectionError after MAX_RETRIES=3, not 5.""" embedder = self._make_embedder() mock_c = AsyncMock() mock_c.__aenter__ = AsyncMock(return_value=mock_c) mock_c.__aexit__ = AsyncMock(return_value=False) mock_c.post = AsyncMock(side_effect=httpx.ConnectError("refused")) with ( patch(_EMBED_PATH, return_value=mock_c), patch("asyncio.sleep", new_callable=AsyncMock), pytest.raises(OllamaConnectionError), ): await embedder.aembed_query("hello") assert mock_c.post.call_count == MAX_RETRIES async def test_timeout_does_not_trigger_429_retry_path(self) -> None: """TimeoutException -> OllamaConnectionError after MAX_RETRIES=3.""" embedder = self._make_embedder() mock_c = AsyncMock() mock_c.__aenter__ = AsyncMock(return_value=mock_c) mock_c.__aexit__ = AsyncMock(return_value=False) mock_c.post = AsyncMock(side_effect=httpx.TimeoutException("timed out")) with ( patch(_EMBED_PATH, return_value=mock_c), patch("asyncio.sleep", new_callable=AsyncMock), pytest.raises(OllamaConnectionError), ): await embedder.aembed_query("hello") assert mock_c.post.call_count == MAX_RETRIES # =========================================================================== # 4. OllamaEmbedder - sync path (embed_query) # =========================================================================== class TestOllamaEmbedderSync: """Tests for the synchronous embed_query method.""" def _make_embedder(self) -> OllamaEmbedder: return OllamaEmbedder(base_url="http://ollama-test:11434") def test_five_consecutive_429s_raise_rate_limit_error(self) -> None: embedder = self._make_embedder() mock_client = MagicMock() mock_client.post.return_value = _429_response() with ( patch.object(embedder, "_get_sync_client", return_value=mock_client), patch("time.sleep"), pytest.raises(RateLimitError) as exc_info, ): embedder.embed_query("hello") assert exc_info.value.provider == "ollama" assert mock_client.post.call_count == MAX_RATE_LIMIT_RETRIES def test_retry_after_header_respected_as_sleep_duration(self) -> None: embedder = self._make_embedder() sleep_calls: list[float] = [] mock_client = MagicMock() mock_client.post.return_value = _429_response(retry_after="9") with ( patch.object(embedder, "_get_sync_client", return_value=mock_client), patch("time.sleep", side_effect=sleep_calls.append), pytest.raises(RateLimitError), ): embedder.embed_query("hello") assert all(s == _RETRY_AFTER_9 for s in sleep_calls) assert len(sleep_calls) == MAX_RATE_LIMIT_RETRIES - 1 def test_partial_retries_then_success(self) -> None: embedder = self._make_embedder() mock_client = MagicMock() mock_client.post.side_effect = [_429_response(), _success_embed_response()] with ( patch.object(embedder, "_get_sync_client", return_value=mock_client), patch("time.sleep"), ): result = embedder.embed_query("hello") assert isinstance(result, list) def test_connect_error_does_not_trigger_rate_limit_retry(self) -> None: embedder = self._make_embedder() mock_client = MagicMock() mock_client.post.side_effect = httpx.ConnectError("refused") with ( patch.object(embedder, "_get_sync_client", return_value=mock_client), patch("time.sleep"), pytest.raises(OllamaConnectionError), ): embedder.embed_query("hello") assert mock_client.post.call_count == MAX_RETRIES # =========================================================================== # 5. OllamaEmbedder - _embed_batch_sync # =========================================================================== class TestOllamaEmbedBatchSync: def _make_embedder(self) -> OllamaEmbedder: return OllamaEmbedder(base_url="http://ollama-test:11434") def test_five_429s_raise_rate_limit_error(self) -> None: embedder = self._make_embedder() mock_client = MagicMock() mock_client.post.return_value = _429_response() with patch("time.sleep"), pytest.raises(RateLimitError): embedder._embed_batch_sync(mock_client, ["a", "b"], batch_index=0) assert mock_client.post.call_count == MAX_RATE_LIMIT_RETRIES def test_connect_error_raises_connection_error_not_rate_limit(self) -> None: embedder = self._make_embedder() mock_client = MagicMock() mock_client.post.side_effect = httpx.ConnectError("refused") with patch("time.sleep"), pytest.raises(OllamaConnectionError): embedder._embed_batch_sync(mock_client, ["a"], batch_index=0) assert mock_client.post.call_count == MAX_RETRIES # =========================================================================== # 6. extraction.py - Anthropic rate-limit retry # =========================================================================== class TestExtractionAnthropicRetry: """Tests for ExtractionService.extract_with_llm Anthropic retry logic. AsyncAnthropic is imported *inside* extract_with_llm, so we patch at the anthropic module level. The client is NOT used as a context manager there. """ async def test_five_rate_limit_errors_raises_rate_limit_error(self) -> None: """RateLimitError raised 5 times -> our RateLimitError propagated.""" svc = ExtractionService() api_exc = _make_anthropic_rl_exc(retry_after="5") mock_client = MagicMock() mock_client.messages.create = AsyncMock(side_effect=api_exc) with ( patch("anthropic.AsyncAnthropic") as mock_cls, patch("roboco.config.settings") as mock_settings, patch("asyncio.sleep", new_callable=AsyncMock), pytest.raises(RateLimitError) as exc_info, ): mock_cls.return_value = mock_client mock_settings.anthropic_api_key = "test-key" await svc.extract_with_llm(_make_extraction_context()) assert exc_info.value.provider == "anthropic" assert mock_client.messages.create.call_count == MAX_RATE_LIMIT_RETRIES async def test_retry_after_header_drives_sleep_duration(self) -> None: """Retry-After: 12 -> asyncio.sleep(12) called on each gap.""" svc = ExtractionService() sleep_calls: list[float] = [] api_exc = _make_anthropic_rl_exc(retry_after="12") mock_client = MagicMock() mock_client.messages.create = AsyncMock(side_effect=api_exc) async def _fake_sleep(secs: float) -> None: sleep_calls.append(secs) with ( patch("anthropic.AsyncAnthropic") as mock_cls, patch("roboco.config.settings") as mock_settings, patch("asyncio.sleep", side_effect=_fake_sleep), pytest.raises(RateLimitError), ): mock_cls.return_value = mock_client mock_settings.anthropic_api_key = "test-key" await svc.extract_with_llm(_make_extraction_context()) assert all(s == _RETRY_AFTER_12 for s in sleep_calls) assert len(sleep_calls) == MAX_RATE_LIMIT_RETRIES - 1 async def test_partial_rate_limit_then_success_returns_result(self) -> None: """Two RateLimitErrors then success -> result returned, no raise.""" svc = ExtractionService() api_exc = _make_anthropic_rl_exc() text_block = MagicMock() text_block.text = "[N,]{type,content,confidence}:\nreasoning,Hello world,0.9" success_response = MagicMock() success_response.content = [text_block] mock_client = MagicMock() mock_client.messages.create = AsyncMock( side_effect=[api_exc, api_exc, success_response] ) raised = False with ( patch("anthropic.AsyncAnthropic") as mock_cls, patch("roboco.config.settings") as mock_settings, patch("asyncio.sleep", new_callable=AsyncMock), ): mock_cls.return_value = mock_client mock_settings.anthropic_api_key = "test-key" try: result = await svc.extract_with_llm(_make_extraction_context()) assert result is not None except RateLimitError: raised = True assert not raised, "RateLimitError raised even though 3rd attempt succeeded" assert mock_client.messages.create.call_count == _CALLS_2RL_1_SUCCESS # =========================================================================== # 7. indexes/base.py - BaseIndexPlugin.ask() LLM 429 retry # =========================================================================== class TestIndexAsk429Retry: """Tests for the LLM call in BaseIndexPlugin.ask().""" async def test_ask_raises_rate_limit_after_five_429s(self) -> None: plugin = _make_journal_plugin() mock_c = _make_async_client_mock(post_return=_429_response()) with ( patch.object( plugin, "search", AsyncMock(return_value=_make_search_outcome()) ), patch("httpx.AsyncClient", return_value=mock_c), patch("asyncio.sleep", new_callable=AsyncMock), pytest.raises(RateLimitError) as exc_info, ): await plugin.ask("what is X") assert exc_info.value.provider == "ollama" assert mock_c.post.call_count == MAX_RATE_LIMIT_RETRIES async def test_ask_partial_retry_then_success(self) -> None: plugin = _make_journal_plugin() success_body = {"choices": [{"message": {"content": "Here is the answer."}}]} success_resp = _make_response(200, success_body) mock_c = AsyncMock() mock_c.__aenter__ = AsyncMock(return_value=mock_c) mock_c.__aexit__ = AsyncMock(return_value=False) mock_c.post = AsyncMock(side_effect=[_429_response(), success_resp]) with ( patch.object( plugin, "search", AsyncMock(return_value=_make_search_outcome()) ), patch("httpx.AsyncClient", return_value=mock_c), patch("asyncio.sleep", new_callable=AsyncMock), ): answer, results = await plugin.ask("what is X") assert answer == "Here is the answer." assert results is not None # =========================================================================== # 8. mentor.py - MentorService._synthesize_answer() 429 retry # =========================================================================== class TestMentorSynthesizeAnswer429: async def test_raises_rate_limit_after_five_429s(self) -> None: mentor = MentorService() mentor._optimal_service = MagicMock() mock_c = _make_async_client_mock(post_return=_429_response()) with ( patch("httpx.AsyncClient", return_value=mock_c), patch("asyncio.sleep", new_callable=AsyncMock), pytest.raises(RateLimitError) as exc_info, ): await mentor._synthesize_answer( question="test", sources=_make_sources(), conversation_context="", agent_profile=None, journal_context=[], ) assert exc_info.value.provider == "ollama" assert mock_c.post.call_count == MAX_RATE_LIMIT_RETRIES async def test_partial_retry_then_success(self) -> None: mentor = MentorService() mentor._optimal_service = MagicMock() success_body = {"choices": [{"message": {"content": "Great answer."}}]} mock_c = AsyncMock() mock_c.__aenter__ = AsyncMock(return_value=mock_c) mock_c.__aexit__ = AsyncMock(return_value=False) mock_c.post = AsyncMock( side_effect=[_429_response(), _make_response(200, success_body)] ) with ( patch("httpx.AsyncClient", return_value=mock_c), patch("asyncio.sleep", new_callable=AsyncMock), ): answer = await mentor._synthesize_answer( question="test", sources=_make_sources(), conversation_context="", agent_profile=None, journal_context=[], ) assert answer == "Great answer." # =========================================================================== # 9. validator.py - ValidatorService._validate_with_llm() 429 retry # =========================================================================== class TestValidatorLLMRetry: async def test_raises_rate_limit_after_five_429s(self) -> None: validator = ValidatorService() validator._optimal_service = MagicMock() validator._llm_available = True mock_c = _make_async_client_mock(post_return=_429_response()) with ( patch("httpx.AsyncClient", return_value=mock_c), patch("asyncio.sleep", new_callable=AsyncMock), pytest.raises(RateLimitError) as exc_info, ): await validator._validate_with_llm( action_type="create_endpoint", context="def foo(): pass", standards=_make_standards(), ) assert exc_info.value.provider == "ollama" assert mock_c.post.call_count == MAX_RATE_LIMIT_RETRIES async def test_partial_retry_then_success(self) -> None: validator = ValidatorService() validator._optimal_service = MagicMock() validator._llm_available = True success_content = '{"violations": [], "summary": "ok"}' success_body = {"choices": [{"message": {"content": success_content}}]} mock_c = AsyncMock() mock_c.__aenter__ = AsyncMock(return_value=mock_c) mock_c.__aexit__ = AsyncMock(return_value=False) mock_c.post = AsyncMock( side_effect=[_429_response(), _make_response(200, success_body)] ) with ( patch("httpx.AsyncClient", return_value=mock_c), patch("asyncio.sleep", new_callable=AsyncMock), ): violations, warnings = await validator._validate_with_llm( action_type="create_endpoint", context="def foo(): pass", standards=_make_standards(), ) assert violations == [] assert warnings == []