wehub-resource-sync
109 lines
4.1 KiB
Python
109 lines
4.1 KiB
Python
"""Tests for Azure OpenAI client."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
import pytest
|
|
|
|
from src.ai.client import AzureOpenAIClient, create_ai_client
|
|
from src.models import AIConfig, AIProvider
|
|
|
|
|
|
def _make_config(**overrides) -> AIConfig:
|
|
defaults = {
|
|
"provider": AIProvider.AZURE,
|
|
"model": "gpt-4o-deployment",
|
|
"api_key_env": "AZURE_OPENAI_API_KEY",
|
|
"azure_endpoint_env": "AZURE_OPENAI_ENDPOINT",
|
|
"api_version": "2024-10-21",
|
|
"temperature": 0.3,
|
|
"max_tokens": 4096,
|
|
}
|
|
defaults.update(overrides)
|
|
return AIConfig(**defaults)
|
|
|
|
|
|
def _mock_response(content: str = '{"ok": true}'):
|
|
return SimpleNamespace(
|
|
choices=[SimpleNamespace(message=SimpleNamespace(content=content))],
|
|
usage=SimpleNamespace(prompt_tokens=10, completion_tokens=5),
|
|
)
|
|
|
|
|
|
class TestAzureOpenAIClientInit:
|
|
def test_creates_instance_with_valid_config(self, monkeypatch):
|
|
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-key")
|
|
monkeypatch.setenv("AZURE_OPENAI_ENDPOINT", "https://example.openai.azure.com")
|
|
|
|
client = AzureOpenAIClient(_make_config())
|
|
|
|
assert client.model == "gpt-4o-deployment"
|
|
assert client.max_tokens == 4096
|
|
|
|
def test_raises_when_api_key_missing(self, monkeypatch):
|
|
monkeypatch.delenv("AZURE_OPENAI_API_KEY", raising=False)
|
|
monkeypatch.setenv("AZURE_OPENAI_ENDPOINT", "https://example.openai.azure.com")
|
|
|
|
with pytest.raises(ValueError, match="Missing API key"):
|
|
AzureOpenAIClient(_make_config())
|
|
|
|
def test_raises_when_endpoint_missing(self, monkeypatch):
|
|
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-key")
|
|
monkeypatch.delenv("AZURE_OPENAI_ENDPOINT", raising=False)
|
|
|
|
with pytest.raises(ValueError, match="Missing Azure endpoint"):
|
|
AzureOpenAIClient(_make_config())
|
|
|
|
|
|
class TestAzureOpenAIClientComplete:
|
|
def test_uses_max_completion_tokens_for_gpt5_prefix(self, monkeypatch):
|
|
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-key")
|
|
monkeypatch.setenv("AZURE_OPENAI_ENDPOINT", "https://example.openai.azure.com")
|
|
client = AzureOpenAIClient(_make_config(model="gpt-5-mini"))
|
|
|
|
with patch.object(
|
|
client.client.chat.completions, "create", new_callable=AsyncMock
|
|
) as mock_create:
|
|
mock_create.return_value = _mock_response()
|
|
asyncio.run(client.complete(system="test", user="hello"))
|
|
|
|
call_kwargs = mock_create.call_args[1]
|
|
assert call_kwargs["max_completion_tokens"] == 4096
|
|
assert "max_tokens" not in call_kwargs
|
|
|
|
def test_retries_with_max_completion_tokens_for_custom_deployment(self, monkeypatch):
|
|
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-key")
|
|
monkeypatch.setenv("AZURE_OPENAI_ENDPOINT", "https://example.openai.azure.com")
|
|
client = AzureOpenAIClient(_make_config(model="prod-gpt5-nano"))
|
|
|
|
with patch.object(
|
|
client.client.chat.completions, "create", new_callable=AsyncMock
|
|
) as mock_create:
|
|
mock_create.side_effect = [
|
|
RuntimeError("Unsupported parameter: 'max_tokens'. Use 'max_completion_tokens' instead."),
|
|
_mock_response(),
|
|
]
|
|
result = asyncio.run(client.complete(system="test", user="hello"))
|
|
|
|
assert result == '{"ok": true}'
|
|
first_call = mock_create.call_args_list[0].kwargs
|
|
second_call = mock_create.call_args_list[1].kwargs
|
|
assert first_call["max_tokens"] == 4096
|
|
assert "max_completion_tokens" not in first_call
|
|
assert second_call["max_completion_tokens"] == 4096
|
|
assert "max_tokens" not in second_call
|
|
assert client._use_max_completion_tokens is True
|
|
|
|
|
|
class TestFactoryFunction:
|
|
def test_creates_azure_client(self, monkeypatch):
|
|
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-key")
|
|
monkeypatch.setenv("AZURE_OPENAI_ENDPOINT", "https://example.openai.azure.com")
|
|
|
|
client = create_ai_client(_make_config())
|
|
|
|
assert isinstance(client, AzureOpenAIClient)
|