mirror of
https://github.com/hoshikawa2/compass_backoffice.git
synced 2026-07-09 22:04:20 +00:00
51 lines
2.2 KiB
Python
51 lines
2.2 KiB
Python
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
|