"""Fixtures for llama.cpp integration tests.""" from collections.abc import AsyncGenerator, Generator from dataclasses import dataclass, field import logging from typing import Any from unittest.mock import AsyncMock, patch import pytest from homeassistant.components import conversation from homeassistant.components.llama_cpp.const import ( CONF_BASE_URL, DEFAULT_BASE_URL, DEFAULT_CONVERSATION_NAME, DOMAIN, ) from homeassistant.config_entries import ConfigSubentryData from homeassistant.const import CONF_API_KEY, CONF_LLM_HASS_API, Platform from homeassistant.core import HomeAssistant from homeassistant.helpers import chat_session, llm from homeassistant.setup import async_setup_component from tests.common import MockConfigEntry _LOGGER = logging.getLogger(__name__) CONFIG_ENTRY_DATA = { CONF_API_KEY: "sk-0000000000000000000", CONF_BASE_URL: DEFAULT_BASE_URL, } ASSIST_OPTIONS = {CONF_LLM_HASS_API: llm.LLM_API_ASSIST} @pytest.fixture(autouse=True) async def setup_home_assistant(hass: HomeAssistant) -> None: """Enable dependencies.""" assert await async_setup_component(hass, "homeassistant", {}) @pytest.fixture(name="platforms") def mock_platforms() -> list[Platform]: """Fixture for platforms loaded by the integration.""" return [] @pytest.fixture(name="setup_integration") async def mock_setup_integration( hass: HomeAssistant, mock_config_entry: MockConfigEntry, platforms: list[Platform], ) -> AsyncGenerator[None]: """Set up the integration.""" with patch(f"homeassistant.components.{DOMAIN}.PLATFORMS", platforms): assert await async_setup_component(hass, DOMAIN, {}) await hass.async_block_till_done() yield @pytest.fixture(name="config_entry_data") def config_entry_data_fixture() -> dict[str, Any]: """Fixture to add data to the config entry.""" return {} @pytest.fixture(name="config_entry_options") def config_entry_options_fixture() -> dict[str, Any]: """Fixture to add options to the config entry.""" return {} @pytest.fixture(name="mock_config_entry") def mock_config_entry_fixture( hass: HomeAssistant, config_entry_data: dict[str, Any], config_entry_options: dict[str, Any], ) -> MockConfigEntry: """Fixture to create a configuration entry.""" config_entry = MockConfigEntry( domain=DOMAIN, title="llama.cpp", data={ **CONFIG_ENTRY_DATA, **config_entry_data, }, version=1, minor_version=1, subentries_data=[ ConfigSubentryData( data={**config_entry_options}, subentry_type="conversation", title=DEFAULT_CONVERSATION_NAME, unique_id=None, ), ], ) config_entry.add_to_hass(hass) return config_entry @dataclass class MockChatLog(conversation.ChatLog): """Mock chat log.""" _mock_tool_results: dict[str, Any] = field(default_factory=dict) def mock_tool_results(self, results: dict[str, Any]) -> None: """Set tool results.""" self._mock_tool_results = results @property def llm_api(self) -> llm.APIInstance | None: """Return LLM API.""" return self._llm_api @llm_api.setter def llm_api(self, value: llm.APIInstance | None) -> None: """Set LLM API.""" self._llm_api = value if not value: return async def async_call_tool(tool_input: llm.ToolInput) -> llm.ToolResult: """Call tool.""" if tool_input.id not in self._mock_tool_results: raise ValueError( f"Tool {tool_input.id} not found ({self._mock_tool_results})" ) return self._mock_tool_results[tool_input.id] self._llm_api.async_call_tool = async_call_tool @pytest.fixture def mock_chat_log(hass: HomeAssistant) -> Generator[conversation.ChatLog]: """Return mock chat logs.""" # pylint: disable-next=contextmanager-generator-missing-cleanup with ( patch( "homeassistant.components.conversation.chat_log.ChatLog", MockChatLog, ), chat_session.async_get_chat_session(hass, "mock-conversation-id") as session, conversation.async_get_chat_log(hass, session) as chat_log, ): yield chat_log @pytest.fixture(autouse=True) def mock_models_list() -> Generator[AsyncMock]: """Initialize integration.""" with patch( "openai.resources.models.AsyncModels.list", new_callable=AsyncMock, ) as mock_list: yield mock_list @pytest.fixture(name="mock_completion", autouse=True) def mock_openai_client_fixture() -> Generator[AsyncMock]: """Fixture to mock the OpenAI client.""" with patch( "openai.resources.chat.completions.AsyncCompletions.create", new_callable=AsyncMock, ) as mock_create: yield mock_create