"""Tests for the OpenAI integration.""" from unittest.mock import patch from httpx import Response from openai import ( APIConnectionError, AuthenticationError, BadRequestError, RateLimitError, ) from openai.types.image import Image from openai.types.images_response import ImagesResponse import pytest from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError, ServiceValidationError from homeassistant.setup import async_setup_component from tests.common import MockConfigEntry @pytest.mark.parametrize( ("service_data", "expected_args"), [ ( {"prompt": "Picture of a dog"}, { "prompt": "Picture of a dog", "size": "1024x1024", "quality": "standard", "style": "vivid", }, ), ( { "prompt": "Picture of a dog", "size": "1024x1792", "quality": "hd", "style": "vivid", }, { "prompt": "Picture of a dog", "size": "1024x1792", "quality": "hd", "style": "vivid", }, ), ( { "prompt": "Picture of a dog", "size": "1792x1024", "quality": "standard", "style": "natural", }, { "prompt": "Picture of a dog", "size": "1792x1024", "quality": "standard", "style": "natural", }, ), ], ) async def test_generate_image_service( hass: HomeAssistant, mock_config_entry: MockConfigEntry, mock_init_component, service_data, expected_args, ) -> None: """Test generate image service.""" service_data["config_entry"] = mock_config_entry.entry_id expected_args["model"] = "dall-e-3" expected_args["response_format"] = "url" expected_args["n"] = 1 with patch( "openai.resources.images.AsyncImages.generate", return_value=ImagesResponse( created=1700000000, data=[ Image( b64_json=None, revised_prompt="A clear and detailed picture of an ordinary canine", url="A", ) ], ), ) as mock_create: response = await hass.services.async_call( "openai_conversation", "generate_image", service_data, blocking=True, return_response=True, ) assert response == { "url": "A", "revised_prompt": "A clear and detailed picture of an ordinary canine", } assert len(mock_create.mock_calls) == 1 assert mock_create.mock_calls[0][2] == expected_args @pytest.mark.usefixtures("mock_init_component") async def test_generate_image_service_error( hass: HomeAssistant, mock_config_entry: MockConfigEntry, ) -> None: """Test generate image service handles errors.""" with ( patch( "openai.resources.images.AsyncImages.generate", side_effect=RateLimitError( response=Response(status_code=None, request=""), body=None, message="Reason", ), ), pytest.raises(HomeAssistantError, match="Error generating image: Reason"), ): await hass.services.async_call( "openai_conversation", "generate_image", { "config_entry": mock_config_entry.entry_id, "prompt": "Image of an epic fail", }, blocking=True, return_response=True, ) async def test_invalid_config_entry( hass: HomeAssistant, mock_config_entry: MockConfigEntry, mock_init_component, ) -> None: """Assert exception when invalid config entry is provided.""" service_data = { "prompt": "Picture of a dog", "config_entry": "invalid_entry", } with pytest.raises( ServiceValidationError, match="Invalid config entry provided. Got invalid_entry" ): await hass.services.async_call( "openai_conversation", "generate_image", service_data, blocking=True, return_response=True, ) @pytest.mark.parametrize( ("side_effect", "error"), [ (APIConnectionError(request=None), "Connection error"), ( AuthenticationError( response=Response(status_code=None, request=""), body=None, message=None ), "Invalid API key", ), ( BadRequestError( response=Response(status_code=None, request=""), body=None, message=None ), "openai_conversation integration not ready yet: None", ), ], ) async def test_init_error( hass: HomeAssistant, mock_config_entry: MockConfigEntry, caplog: pytest.LogCaptureFixture, side_effect, error, ) -> None: """Test initialization errors.""" with patch( "openai.resources.models.AsyncModels.list", side_effect=side_effect, ): assert await async_setup_component(hass, "openai_conversation", {}) await hass.async_block_till_done() assert error in caplog.text