first commit

tests_original_develop/unit/test_llm_retry.py

@@ -0,0 +1,50 @@
+import pytest
+from unittest.mock import patch, MagicMock
+import oci
+from src.providers.llm_provider import chat_llm_with_usage, classification_llm
+
+@patch("oci.generative_ai_inference.GenerativeAiInferenceClient.chat")
+def test_chat_llm_retry_on_timeout(mock_chat):
+    """
+    Testa se o chat_llm_with_usage executa retries quando recebe um ServiceError (504).
+    Simula 2 falhas seguidas e sucesso na 3ª tentativa.
+    """
+    # ── Montagem do Cenário ──────────────────────────────────────────────────
+    # Criamos o erro que o retry deve capturar
+    timeout_error = oci.exceptions.ServiceError(
+        status=504,
+        code="GatewayTimeout",
+        message="Request timed out",
+        headers={},
+        operation_name="Chat"
+    )
+    
+    # Simula resposta de sucesso na 3ª chamada
+    mock_response = MagicMock()
+    chat_response = mock_response.data.chat_response
+    
+    # Configuramos os valores de retorno para serem tipos primitivos (não Mocks)
+    chat_response.choices = [
+        MagicMock(
+            message=MagicMock(content=[MagicMock(text="Sucesso após retry!")]),
+            finish_reason="stop"
+        )
+    ]
+    chat_response.model = "cohere.command-r-plus"
+    chat_response.usage.prompt_tokens = 10
+    chat_response.usage.completion_tokens = 10
+    chat_response.usage.total_tokens = 20
+    
+    # Configuramos os side_effects: Falha, Falha, Sucesso
+    mock_chat.side_effect = [timeout_error, timeout_error, mock_response]
+
+    # ── Execução ─────────────────────────────────────────────────────────────
+    prompt = "Teste de retry"
+    result = chat_llm_with_usage(classification_llm, prompt)
+
+    # ── Verificação ──────────────────────────────────────────────────────────
+    # O resultado final deve ser o da 3ª chamada
+    assert result.content == "Sucesso após retry!"
+    
+    # Verificamos se o client.chat foi chamado exatamente 3 vezes
+    assert mock_chat.call_count == 3