mirror of
https://github.com/home-assistant/core.git
synced 2025-11-10 03:19:34 +00:00
1000 lines
34 KiB
Python
1000 lines
34 KiB
Python
"""Tests for the Ollama integration."""
|
|
|
|
from typing import Any
|
|
from unittest.mock import patch
|
|
|
|
from httpx import ConnectError
|
|
import pytest
|
|
|
|
from homeassistant.components import ollama
|
|
from homeassistant.components.ollama.const import DOMAIN
|
|
from homeassistant.config_entries import ConfigEntryDisabler, ConfigSubentryData
|
|
from homeassistant.const import CONF_URL
|
|
from homeassistant.core import HomeAssistant
|
|
from homeassistant.helpers import device_registry as dr, entity_registry as er, llm
|
|
from homeassistant.helpers.device_registry import DeviceEntryDisabler
|
|
from homeassistant.helpers.entity_registry import RegistryEntryDisabler
|
|
from homeassistant.setup import async_setup_component
|
|
|
|
from . import TEST_OPTIONS
|
|
|
|
from tests.common import MockConfigEntry
|
|
|
|
V1_TEST_USER_DATA = {
|
|
ollama.CONF_URL: "http://localhost:11434",
|
|
ollama.CONF_MODEL: "test_model:latest",
|
|
}
|
|
|
|
V1_TEST_OPTIONS = {
|
|
ollama.CONF_PROMPT: llm.DEFAULT_INSTRUCTIONS_PROMPT,
|
|
ollama.CONF_MAX_HISTORY: 2,
|
|
}
|
|
|
|
V21_TEST_USER_DATA = V1_TEST_USER_DATA
|
|
V21_TEST_OPTIONS = V1_TEST_OPTIONS
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("side_effect", "error"),
|
|
[
|
|
(ConnectError(message="Connect error"), "Connect error"),
|
|
(RuntimeError("Runtime error"), "Runtime error"),
|
|
],
|
|
)
|
|
async def test_init_error(
|
|
hass: HomeAssistant,
|
|
mock_config_entry: MockConfigEntry,
|
|
caplog: pytest.LogCaptureFixture,
|
|
side_effect,
|
|
error,
|
|
) -> None:
|
|
"""Test initialization errors."""
|
|
with patch(
|
|
"ollama.AsyncClient.list",
|
|
side_effect=side_effect,
|
|
):
|
|
assert await async_setup_component(hass, ollama.DOMAIN, {})
|
|
await hass.async_block_till_done()
|
|
assert error in caplog.text
|
|
|
|
|
|
async def test_migration_from_v1(
|
|
hass: HomeAssistant,
|
|
device_registry: dr.DeviceRegistry,
|
|
entity_registry: er.EntityRegistry,
|
|
) -> None:
|
|
"""Test migration from version 1."""
|
|
# Create a v1 config entry with conversation options and an entity
|
|
mock_config_entry = MockConfigEntry(
|
|
domain=DOMAIN,
|
|
data=V1_TEST_USER_DATA,
|
|
options=V1_TEST_OPTIONS,
|
|
version=1,
|
|
title="llama-3.2-8b",
|
|
)
|
|
mock_config_entry.add_to_hass(hass)
|
|
|
|
device = device_registry.async_get_or_create(
|
|
config_entry_id=mock_config_entry.entry_id,
|
|
identifiers={(DOMAIN, mock_config_entry.entry_id)},
|
|
name=mock_config_entry.title,
|
|
manufacturer="Ollama",
|
|
model="Ollama",
|
|
entry_type=dr.DeviceEntryType.SERVICE,
|
|
)
|
|
entity = entity_registry.async_get_or_create(
|
|
"conversation",
|
|
DOMAIN,
|
|
mock_config_entry.entry_id,
|
|
config_entry=mock_config_entry,
|
|
device_id=device.id,
|
|
suggested_object_id="llama_3_2_8b",
|
|
)
|
|
|
|
# Run migration
|
|
with patch(
|
|
"homeassistant.components.ollama.async_setup_entry",
|
|
return_value=True,
|
|
):
|
|
await hass.config_entries.async_setup(mock_config_entry.entry_id)
|
|
await hass.async_block_till_done()
|
|
|
|
assert mock_config_entry.version == 3
|
|
assert mock_config_entry.minor_version == 3
|
|
# After migration, parent entry should only have URL
|
|
assert mock_config_entry.data == {ollama.CONF_URL: "http://localhost:11434"}
|
|
assert mock_config_entry.options == {}
|
|
|
|
assert len(mock_config_entry.subentries) == 2
|
|
|
|
subentry = next(
|
|
iter(
|
|
entry
|
|
for entry in mock_config_entry.subentries.values()
|
|
if entry.subentry_type == "conversation"
|
|
)
|
|
)
|
|
assert subentry.unique_id is None
|
|
assert subentry.title == "llama-3.2-8b"
|
|
assert subentry.subentry_type == "conversation"
|
|
# Subentry should now include the model from the original options
|
|
expected_subentry_data = TEST_OPTIONS.copy()
|
|
assert subentry.data == expected_subentry_data
|
|
|
|
# Find the AI Task subentry
|
|
ai_task_subentry = next(
|
|
iter(
|
|
entry
|
|
for entry in mock_config_entry.subentries.values()
|
|
if entry.subentry_type == "ai_task_data"
|
|
)
|
|
)
|
|
assert ai_task_subentry.unique_id is None
|
|
assert ai_task_subentry.title == "Ollama AI Task"
|
|
assert ai_task_subentry.subentry_type == "ai_task_data"
|
|
|
|
migrated_entity = entity_registry.async_get(entity.entity_id)
|
|
assert migrated_entity is not None
|
|
assert migrated_entity.config_entry_id == mock_config_entry.entry_id
|
|
assert migrated_entity.config_subentry_id == subentry.subentry_id
|
|
assert migrated_entity.unique_id == subentry.subentry_id
|
|
|
|
# Check device migration
|
|
assert not device_registry.async_get_device(
|
|
identifiers={(DOMAIN, mock_config_entry.entry_id)}
|
|
)
|
|
assert (
|
|
migrated_device := device_registry.async_get_device(
|
|
identifiers={(DOMAIN, subentry.subentry_id)}
|
|
)
|
|
)
|
|
assert migrated_device.identifiers == {(DOMAIN, subentry.subentry_id)}
|
|
assert migrated_device.id == device.id
|
|
assert migrated_device.config_entries == {mock_config_entry.entry_id}
|
|
assert migrated_device.config_entries_subentries == {
|
|
mock_config_entry.entry_id: {subentry.subentry_id}
|
|
}
|
|
|
|
|
|
async def test_migration_from_v1_with_multiple_urls(
|
|
hass: HomeAssistant,
|
|
device_registry: dr.DeviceRegistry,
|
|
entity_registry: er.EntityRegistry,
|
|
) -> None:
|
|
"""Test migration from version 1 with different URLs."""
|
|
# Create two v1 config entries with different URLs
|
|
mock_config_entry = MockConfigEntry(
|
|
domain=DOMAIN,
|
|
data={"url": "http://localhost:11434", "model": "llama3.2:latest"},
|
|
options=V1_TEST_OPTIONS,
|
|
version=1,
|
|
title="Ollama 1",
|
|
)
|
|
mock_config_entry.add_to_hass(hass)
|
|
mock_config_entry_2 = MockConfigEntry(
|
|
domain=DOMAIN,
|
|
data={"url": "http://localhost:11435", "model": "llama3.2:latest"},
|
|
options=V1_TEST_OPTIONS,
|
|
version=1,
|
|
title="Ollama 2",
|
|
)
|
|
mock_config_entry_2.add_to_hass(hass)
|
|
|
|
device = device_registry.async_get_or_create(
|
|
config_entry_id=mock_config_entry.entry_id,
|
|
identifiers={(DOMAIN, mock_config_entry.entry_id)},
|
|
name=mock_config_entry.title,
|
|
manufacturer="Ollama",
|
|
model="Ollama 1",
|
|
entry_type=dr.DeviceEntryType.SERVICE,
|
|
)
|
|
entity_registry.async_get_or_create(
|
|
"conversation",
|
|
DOMAIN,
|
|
mock_config_entry.entry_id,
|
|
config_entry=mock_config_entry,
|
|
device_id=device.id,
|
|
suggested_object_id="ollama_1",
|
|
)
|
|
|
|
device_2 = device_registry.async_get_or_create(
|
|
config_entry_id=mock_config_entry_2.entry_id,
|
|
identifiers={(DOMAIN, mock_config_entry_2.entry_id)},
|
|
name=mock_config_entry_2.title,
|
|
manufacturer="Ollama",
|
|
model="Ollama 2",
|
|
entry_type=dr.DeviceEntryType.SERVICE,
|
|
)
|
|
entity_registry.async_get_or_create(
|
|
"conversation",
|
|
DOMAIN,
|
|
mock_config_entry_2.entry_id,
|
|
config_entry=mock_config_entry_2,
|
|
device_id=device_2.id,
|
|
suggested_object_id="ollama_2",
|
|
)
|
|
|
|
# Run migration
|
|
with patch(
|
|
"homeassistant.components.ollama.async_setup_entry",
|
|
return_value=True,
|
|
):
|
|
await hass.config_entries.async_setup(mock_config_entry.entry_id)
|
|
await hass.async_block_till_done()
|
|
|
|
entries = hass.config_entries.async_entries(DOMAIN)
|
|
assert len(entries) == 2
|
|
|
|
for idx, entry in enumerate(entries):
|
|
assert entry.version == 3
|
|
assert entry.minor_version == 3
|
|
assert not entry.options
|
|
assert len(entry.subentries) == 2
|
|
|
|
subentry = next(
|
|
iter(
|
|
subentry
|
|
for subentry in entry.subentries.values()
|
|
if subentry.subentry_type == "conversation"
|
|
)
|
|
)
|
|
assert subentry.subentry_type == "conversation"
|
|
# Subentry should include the model along with the original options
|
|
expected_subentry_data = TEST_OPTIONS.copy()
|
|
expected_subentry_data["model"] = "llama3.2:latest"
|
|
assert subentry.data == expected_subentry_data
|
|
assert subentry.title == f"Ollama {idx + 1}"
|
|
|
|
# Find the AI Task subentry
|
|
ai_task_subentry = next(
|
|
iter(
|
|
subentry
|
|
for subentry in entry.subentries.values()
|
|
if subentry.subentry_type == "ai_task_data"
|
|
)
|
|
)
|
|
assert ai_task_subentry.subentry_type == "ai_task_data"
|
|
assert ai_task_subentry.title == "Ollama AI Task"
|
|
|
|
dev = device_registry.async_get_device(
|
|
identifiers={(DOMAIN, list(entry.subentries.values())[0].subentry_id)}
|
|
)
|
|
assert dev is not None
|
|
assert dev.config_entries == {entry.entry_id}
|
|
assert dev.config_entries_subentries == {entry.entry_id: {subentry.subentry_id}}
|
|
|
|
|
|
async def test_migration_from_v1_with_same_urls(
|
|
hass: HomeAssistant,
|
|
device_registry: dr.DeviceRegistry,
|
|
entity_registry: er.EntityRegistry,
|
|
) -> None:
|
|
"""Test migration from version 1 with same URLs consolidates entries."""
|
|
# Create two v1 config entries with the same URL
|
|
mock_config_entry = MockConfigEntry(
|
|
domain=DOMAIN,
|
|
data={"url": "http://localhost:11434", "model": "llama3.2:latest"},
|
|
options=V1_TEST_OPTIONS,
|
|
version=1,
|
|
title="Ollama",
|
|
)
|
|
mock_config_entry.add_to_hass(hass)
|
|
mock_config_entry_2 = MockConfigEntry(
|
|
domain=DOMAIN,
|
|
data={"url": "http://localhost:11434", "model": "llama3.2:latest"}, # Same URL
|
|
options=V1_TEST_OPTIONS,
|
|
version=1,
|
|
title="Ollama 2",
|
|
)
|
|
mock_config_entry_2.add_to_hass(hass)
|
|
|
|
device = device_registry.async_get_or_create(
|
|
config_entry_id=mock_config_entry.entry_id,
|
|
identifiers={(DOMAIN, mock_config_entry.entry_id)},
|
|
name=mock_config_entry.title,
|
|
manufacturer="Ollama",
|
|
model="Ollama",
|
|
entry_type=dr.DeviceEntryType.SERVICE,
|
|
)
|
|
entity_registry.async_get_or_create(
|
|
"conversation",
|
|
DOMAIN,
|
|
mock_config_entry.entry_id,
|
|
config_entry=mock_config_entry,
|
|
device_id=device.id,
|
|
suggested_object_id="ollama",
|
|
)
|
|
|
|
device_2 = device_registry.async_get_or_create(
|
|
config_entry_id=mock_config_entry_2.entry_id,
|
|
identifiers={(DOMAIN, mock_config_entry_2.entry_id)},
|
|
name=mock_config_entry_2.title,
|
|
manufacturer="Ollama",
|
|
model="Ollama",
|
|
entry_type=dr.DeviceEntryType.SERVICE,
|
|
)
|
|
entity_registry.async_get_or_create(
|
|
"conversation",
|
|
DOMAIN,
|
|
mock_config_entry_2.entry_id,
|
|
config_entry=mock_config_entry_2,
|
|
device_id=device_2.id,
|
|
suggested_object_id="ollama_2",
|
|
)
|
|
|
|
# Run migration
|
|
with patch(
|
|
"homeassistant.components.ollama.async_setup_entry",
|
|
return_value=True,
|
|
):
|
|
await hass.config_entries.async_setup(mock_config_entry.entry_id)
|
|
await hass.async_block_till_done()
|
|
|
|
# Should have only one entry left (consolidated)
|
|
entries = hass.config_entries.async_entries(DOMAIN)
|
|
assert len(entries) == 1
|
|
|
|
entry = entries[0]
|
|
assert entry.version == 3
|
|
assert entry.minor_version == 3
|
|
assert not entry.options
|
|
# Two conversation subentries from the two original entries and 1 aitask subentry
|
|
assert len(entry.subentries) == 3
|
|
|
|
# Check both subentries exist with correct data
|
|
subentries = list(entry.subentries.values())
|
|
titles = [sub.title for sub in subentries]
|
|
assert "Ollama" in titles
|
|
assert "Ollama 2" in titles
|
|
|
|
conversation_subentries = [
|
|
subentry for subentry in subentries if subentry.subentry_type == "conversation"
|
|
]
|
|
assert len(conversation_subentries) == 2
|
|
for subentry in conversation_subentries:
|
|
assert subentry.subentry_type == "conversation"
|
|
# Subentry should include the model along with the original options
|
|
expected_subentry_data = TEST_OPTIONS.copy()
|
|
expected_subentry_data["model"] = "llama3.2:latest"
|
|
assert subentry.data == expected_subentry_data
|
|
|
|
# Check devices were migrated correctly
|
|
dev = device_registry.async_get_device(
|
|
identifiers={(DOMAIN, subentry.subentry_id)}
|
|
)
|
|
assert dev is not None
|
|
assert dev.config_entries == {mock_config_entry.entry_id}
|
|
assert dev.config_entries_subentries == {
|
|
mock_config_entry.entry_id: {subentry.subentry_id}
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
(
|
|
"config_entry_disabled_by",
|
|
"device_disabled_by",
|
|
"entity_disabled_by",
|
|
"merged_config_entry_disabled_by",
|
|
"conversation_subentry_data",
|
|
"main_config_entry",
|
|
),
|
|
[
|
|
(
|
|
[ConfigEntryDisabler.USER, None],
|
|
[DeviceEntryDisabler.CONFIG_ENTRY, None],
|
|
[RegistryEntryDisabler.CONFIG_ENTRY, None],
|
|
None,
|
|
[
|
|
{
|
|
"conversation_entity_id": "conversation.ollama_2",
|
|
"device_disabled_by": None,
|
|
"entity_disabled_by": None,
|
|
"device": 1,
|
|
},
|
|
{
|
|
"conversation_entity_id": "conversation.ollama",
|
|
"device_disabled_by": DeviceEntryDisabler.USER,
|
|
"entity_disabled_by": RegistryEntryDisabler.DEVICE,
|
|
"device": 0,
|
|
},
|
|
],
|
|
1,
|
|
),
|
|
(
|
|
[None, ConfigEntryDisabler.USER],
|
|
[None, DeviceEntryDisabler.CONFIG_ENTRY],
|
|
[None, RegistryEntryDisabler.CONFIG_ENTRY],
|
|
None,
|
|
[
|
|
{
|
|
"conversation_entity_id": "conversation.ollama",
|
|
"device_disabled_by": None,
|
|
"entity_disabled_by": None,
|
|
"device": 0,
|
|
},
|
|
{
|
|
"conversation_entity_id": "conversation.ollama_2",
|
|
"device_disabled_by": DeviceEntryDisabler.USER,
|
|
"entity_disabled_by": RegistryEntryDisabler.DEVICE,
|
|
"device": 1,
|
|
},
|
|
],
|
|
0,
|
|
),
|
|
(
|
|
[ConfigEntryDisabler.USER, ConfigEntryDisabler.USER],
|
|
[DeviceEntryDisabler.CONFIG_ENTRY, DeviceEntryDisabler.CONFIG_ENTRY],
|
|
[RegistryEntryDisabler.CONFIG_ENTRY, RegistryEntryDisabler.CONFIG_ENTRY],
|
|
ConfigEntryDisabler.USER,
|
|
[
|
|
{
|
|
"conversation_entity_id": "conversation.ollama",
|
|
"device_disabled_by": DeviceEntryDisabler.CONFIG_ENTRY,
|
|
"entity_disabled_by": RegistryEntryDisabler.CONFIG_ENTRY,
|
|
"device": 0,
|
|
},
|
|
{
|
|
"conversation_entity_id": "conversation.ollama_2",
|
|
"device_disabled_by": DeviceEntryDisabler.CONFIG_ENTRY,
|
|
"entity_disabled_by": RegistryEntryDisabler.CONFIG_ENTRY,
|
|
"device": 1,
|
|
},
|
|
],
|
|
0,
|
|
),
|
|
],
|
|
)
|
|
async def test_migration_from_v1_disabled(
|
|
hass: HomeAssistant,
|
|
device_registry: dr.DeviceRegistry,
|
|
entity_registry: er.EntityRegistry,
|
|
config_entry_disabled_by: list[ConfigEntryDisabler | None],
|
|
device_disabled_by: list[DeviceEntryDisabler | None],
|
|
entity_disabled_by: list[RegistryEntryDisabler | None],
|
|
merged_config_entry_disabled_by: ConfigEntryDisabler | None,
|
|
conversation_subentry_data: list[dict[str, Any]],
|
|
main_config_entry: int,
|
|
) -> None:
|
|
"""Test migration where the config entries are disabled."""
|
|
# Create a v1 config entry with conversation options and an entity
|
|
mock_config_entry = MockConfigEntry(
|
|
domain=DOMAIN,
|
|
data={"url": "http://localhost:11434", "model": "llama3.2:latest"},
|
|
options=V1_TEST_OPTIONS,
|
|
version=1,
|
|
title="Ollama",
|
|
disabled_by=config_entry_disabled_by[0],
|
|
)
|
|
mock_config_entry.add_to_hass(hass)
|
|
mock_config_entry_2 = MockConfigEntry(
|
|
domain=DOMAIN,
|
|
data={"url": "http://localhost:11434", "model": "llama3.2:latest"},
|
|
options=V1_TEST_OPTIONS,
|
|
version=1,
|
|
title="Ollama 2",
|
|
disabled_by=config_entry_disabled_by[1],
|
|
)
|
|
mock_config_entry_2.add_to_hass(hass)
|
|
mock_config_entries = [mock_config_entry, mock_config_entry_2]
|
|
|
|
device_1 = device_registry.async_get_or_create(
|
|
config_entry_id=mock_config_entry.entry_id,
|
|
identifiers={(DOMAIN, mock_config_entry.entry_id)},
|
|
name=mock_config_entry.title,
|
|
manufacturer="Ollama",
|
|
model="Ollama",
|
|
entry_type=dr.DeviceEntryType.SERVICE,
|
|
disabled_by=device_disabled_by[0],
|
|
)
|
|
entity_registry.async_get_or_create(
|
|
"conversation",
|
|
DOMAIN,
|
|
mock_config_entry.entry_id,
|
|
config_entry=mock_config_entry,
|
|
device_id=device_1.id,
|
|
suggested_object_id="ollama",
|
|
disabled_by=entity_disabled_by[0],
|
|
)
|
|
|
|
device_2 = device_registry.async_get_or_create(
|
|
config_entry_id=mock_config_entry_2.entry_id,
|
|
identifiers={(DOMAIN, mock_config_entry_2.entry_id)},
|
|
name=mock_config_entry_2.title,
|
|
manufacturer="Ollama",
|
|
model="Ollama",
|
|
entry_type=dr.DeviceEntryType.SERVICE,
|
|
disabled_by=device_disabled_by[1],
|
|
)
|
|
entity_registry.async_get_or_create(
|
|
"conversation",
|
|
DOMAIN,
|
|
mock_config_entry_2.entry_id,
|
|
config_entry=mock_config_entry_2,
|
|
device_id=device_2.id,
|
|
suggested_object_id="ollama_2",
|
|
disabled_by=entity_disabled_by[1],
|
|
)
|
|
|
|
devices = [device_1, device_2]
|
|
|
|
# Run migration
|
|
with patch(
|
|
"homeassistant.components.ollama.async_setup_entry",
|
|
return_value=True,
|
|
):
|
|
await hass.config_entries.async_setup(mock_config_entry.entry_id)
|
|
await hass.async_block_till_done()
|
|
|
|
entries = hass.config_entries.async_entries(DOMAIN)
|
|
assert len(entries) == 1
|
|
entry = entries[0]
|
|
assert entry.disabled_by is merged_config_entry_disabled_by
|
|
assert entry.version == 3
|
|
assert entry.minor_version == 3
|
|
assert not entry.options
|
|
assert entry.title == "Ollama"
|
|
assert len(entry.subentries) == 3
|
|
conversation_subentries = [
|
|
subentry
|
|
for subentry in entry.subentries.values()
|
|
if subentry.subentry_type == "conversation"
|
|
]
|
|
assert len(conversation_subentries) == 2
|
|
for subentry in conversation_subentries:
|
|
assert subentry.subentry_type == "conversation"
|
|
assert subentry.data == {"model": "llama3.2:latest", **V1_TEST_OPTIONS}
|
|
assert "Ollama" in subentry.title
|
|
ai_task_subentries = [
|
|
subentry
|
|
for subentry in entry.subentries.values()
|
|
if subentry.subentry_type == "ai_task_data"
|
|
]
|
|
assert len(ai_task_subentries) == 1
|
|
assert ai_task_subentries[0].data == {"model": "llama3.2:latest"}
|
|
assert ai_task_subentries[0].title == "Ollama AI Task"
|
|
|
|
assert not device_registry.async_get_device(
|
|
identifiers={(DOMAIN, mock_config_entry.entry_id)}
|
|
)
|
|
assert not device_registry.async_get_device(
|
|
identifiers={(DOMAIN, mock_config_entry_2.entry_id)}
|
|
)
|
|
|
|
for idx, subentry in enumerate(conversation_subentries):
|
|
subentry_data = conversation_subentry_data[idx]
|
|
entity = entity_registry.async_get(subentry_data["conversation_entity_id"])
|
|
assert entity.unique_id == subentry.subentry_id
|
|
assert entity.config_subentry_id == subentry.subentry_id
|
|
assert entity.config_entry_id == entry.entry_id
|
|
assert entity.disabled_by is subentry_data["entity_disabled_by"]
|
|
|
|
assert (
|
|
device := device_registry.async_get_device(
|
|
identifiers={(DOMAIN, subentry.subentry_id)}
|
|
)
|
|
)
|
|
assert device.identifiers == {(DOMAIN, subentry.subentry_id)}
|
|
assert device.id == devices[subentry_data["device"]].id
|
|
assert device.config_entries == {
|
|
mock_config_entries[main_config_entry].entry_id
|
|
}
|
|
assert device.config_entries_subentries == {
|
|
mock_config_entries[main_config_entry].entry_id: {subentry.subentry_id}
|
|
}
|
|
assert device.disabled_by is subentry_data["device_disabled_by"]
|
|
|
|
|
|
async def test_migration_from_v2_1(
|
|
hass: HomeAssistant,
|
|
device_registry: dr.DeviceRegistry,
|
|
entity_registry: er.EntityRegistry,
|
|
) -> None:
|
|
"""Test migration from version 2.1.
|
|
|
|
This tests we clean up the broken migration in Home Assistant Core
|
|
2025.7.0b0-2025.7.0b1:
|
|
- Fix device registry (Fixed in Home Assistant Core 2025.7.0b2)
|
|
"""
|
|
# Create a v2.1 config entry with 2 subentries, devices and entities
|
|
mock_config_entry = MockConfigEntry(
|
|
domain=DOMAIN,
|
|
data=V21_TEST_USER_DATA,
|
|
entry_id="mock_entry_id",
|
|
version=2,
|
|
minor_version=1,
|
|
subentries_data=[
|
|
ConfigSubentryData(
|
|
data=V21_TEST_OPTIONS,
|
|
subentry_id="mock_id_1",
|
|
subentry_type="conversation",
|
|
title="Ollama",
|
|
unique_id=None,
|
|
),
|
|
ConfigSubentryData(
|
|
data=V21_TEST_OPTIONS,
|
|
subentry_id="mock_id_2",
|
|
subentry_type="conversation",
|
|
title="Ollama 2",
|
|
unique_id=None,
|
|
),
|
|
],
|
|
title="Ollama",
|
|
)
|
|
mock_config_entry.add_to_hass(hass)
|
|
|
|
device_1 = device_registry.async_get_or_create(
|
|
config_entry_id=mock_config_entry.entry_id,
|
|
config_subentry_id="mock_id_1",
|
|
identifiers={(DOMAIN, "mock_id_1")},
|
|
name="Ollama",
|
|
manufacturer="Ollama",
|
|
model="Ollama",
|
|
entry_type=dr.DeviceEntryType.SERVICE,
|
|
)
|
|
device_1 = device_registry.async_update_device(
|
|
device_1.id, add_config_entry_id="mock_entry_id", add_config_subentry_id=None
|
|
)
|
|
assert device_1.config_entries_subentries == {"mock_entry_id": {None, "mock_id_1"}}
|
|
entity_registry.async_get_or_create(
|
|
"conversation",
|
|
DOMAIN,
|
|
"mock_id_1",
|
|
config_entry=mock_config_entry,
|
|
config_subentry_id="mock_id_1",
|
|
device_id=device_1.id,
|
|
suggested_object_id="ollama",
|
|
)
|
|
|
|
device_2 = device_registry.async_get_or_create(
|
|
config_entry_id=mock_config_entry.entry_id,
|
|
config_subentry_id="mock_id_2",
|
|
identifiers={(DOMAIN, "mock_id_2")},
|
|
name="Ollama 2",
|
|
manufacturer="Ollama",
|
|
model="Ollama",
|
|
entry_type=dr.DeviceEntryType.SERVICE,
|
|
)
|
|
entity_registry.async_get_or_create(
|
|
"conversation",
|
|
DOMAIN,
|
|
"mock_id_2",
|
|
config_entry=mock_config_entry,
|
|
config_subentry_id="mock_id_2",
|
|
device_id=device_2.id,
|
|
suggested_object_id="ollama_2",
|
|
)
|
|
|
|
# Run migration
|
|
with patch(
|
|
"homeassistant.components.ollama.async_setup_entry",
|
|
return_value=True,
|
|
):
|
|
await hass.config_entries.async_setup(mock_config_entry.entry_id)
|
|
await hass.async_block_till_done()
|
|
|
|
entries = hass.config_entries.async_entries(DOMAIN)
|
|
assert len(entries) == 1
|
|
entry = entries[0]
|
|
assert entry.version == 3
|
|
assert entry.minor_version == 3
|
|
assert not entry.options
|
|
assert entry.title == "Ollama"
|
|
assert len(entry.subentries) == 3
|
|
conversation_subentries = [
|
|
subentry
|
|
for subentry in entry.subentries.values()
|
|
if subentry.subentry_type == "conversation"
|
|
]
|
|
assert len(conversation_subentries) == 2
|
|
for subentry in conversation_subentries:
|
|
assert subentry.subentry_type == "conversation"
|
|
# Since TEST_USER_DATA no longer has a model, subentry data should be TEST_OPTIONS
|
|
assert subentry.data == TEST_OPTIONS
|
|
assert "Ollama" in subentry.title
|
|
|
|
subentry = conversation_subentries[0]
|
|
|
|
entity = entity_registry.async_get("conversation.ollama")
|
|
assert entity.unique_id == subentry.subentry_id
|
|
assert entity.config_subentry_id == subentry.subentry_id
|
|
assert entity.config_entry_id == entry.entry_id
|
|
|
|
assert not device_registry.async_get_device(
|
|
identifiers={(DOMAIN, mock_config_entry.entry_id)}
|
|
)
|
|
assert (
|
|
device := device_registry.async_get_device(
|
|
identifiers={(DOMAIN, subentry.subentry_id)}
|
|
)
|
|
)
|
|
assert device.identifiers == {(DOMAIN, subentry.subentry_id)}
|
|
assert device.id == device_1.id
|
|
assert device.config_entries == {mock_config_entry.entry_id}
|
|
assert device.config_entries_subentries == {
|
|
mock_config_entry.entry_id: {subentry.subentry_id}
|
|
}
|
|
|
|
subentry = conversation_subentries[1]
|
|
|
|
entity = entity_registry.async_get("conversation.ollama_2")
|
|
assert entity.unique_id == subentry.subentry_id
|
|
assert entity.config_subentry_id == subentry.subentry_id
|
|
assert entity.config_entry_id == entry.entry_id
|
|
assert not device_registry.async_get_device(
|
|
identifiers={(DOMAIN, mock_config_entry.entry_id)}
|
|
)
|
|
assert (
|
|
device := device_registry.async_get_device(
|
|
identifiers={(DOMAIN, subentry.subentry_id)}
|
|
)
|
|
)
|
|
assert device.identifiers == {(DOMAIN, subentry.subentry_id)}
|
|
assert device.id == device_2.id
|
|
assert device.config_entries == {mock_config_entry.entry_id}
|
|
assert device.config_entries_subentries == {
|
|
mock_config_entry.entry_id: {subentry.subentry_id}
|
|
}
|
|
|
|
|
|
async def test_migration_from_v2_2(hass: HomeAssistant) -> None:
|
|
"""Test migration from version 2.2."""
|
|
subentry_data = ConfigSubentryData(
|
|
data=V21_TEST_USER_DATA,
|
|
subentry_type="conversation",
|
|
title="Test Conversation",
|
|
unique_id=None,
|
|
)
|
|
|
|
mock_config_entry = MockConfigEntry(
|
|
domain=DOMAIN,
|
|
data={
|
|
ollama.CONF_URL: "http://localhost:11434",
|
|
ollama.CONF_MODEL: "test_model:latest", # Model still in main data
|
|
},
|
|
version=2,
|
|
minor_version=2,
|
|
subentries_data=[subentry_data],
|
|
)
|
|
mock_config_entry.add_to_hass(hass)
|
|
|
|
with patch(
|
|
"homeassistant.components.ollama.async_setup_entry",
|
|
return_value=True,
|
|
):
|
|
await hass.config_entries.async_setup(mock_config_entry.entry_id)
|
|
|
|
# Check migration to v3.1
|
|
assert mock_config_entry.version == 3
|
|
assert mock_config_entry.minor_version == 3
|
|
|
|
# Check that model was moved from main data to subentry
|
|
assert mock_config_entry.data == {ollama.CONF_URL: "http://localhost:11434"}
|
|
assert len(mock_config_entry.subentries) == 2
|
|
|
|
subentry = next(iter(mock_config_entry.subentries.values()))
|
|
assert subentry.data == {
|
|
**V21_TEST_USER_DATA,
|
|
ollama.CONF_MODEL: "test_model:latest",
|
|
}
|
|
|
|
|
|
async def test_migration_from_v3_1_without_subentry(hass: HomeAssistant) -> None:
|
|
"""Test migration from version 3.1 where there is no existing subentry.
|
|
|
|
This exercises the code path where the model is not moved to a subentry
|
|
because the subentry does not exist, which is a scenario that can happen
|
|
if the user created the config entry without adding a subentry, or
|
|
if the user manually removed the subentry after the migration to v3.1.
|
|
"""
|
|
mock_config_entry = MockConfigEntry(
|
|
domain=DOMAIN,
|
|
data={
|
|
ollama.CONF_MODEL: "test_model:latest",
|
|
},
|
|
version=3,
|
|
minor_version=1,
|
|
)
|
|
mock_config_entry.add_to_hass(hass)
|
|
|
|
with patch(
|
|
"homeassistant.components.ollama.async_setup_entry",
|
|
return_value=True,
|
|
):
|
|
await hass.config_entries.async_setup(mock_config_entry.entry_id)
|
|
|
|
assert mock_config_entry.version == 3
|
|
assert mock_config_entry.minor_version == 3
|
|
|
|
assert next(iter(mock_config_entry.subentries.values()), None) is None
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
(
|
|
"config_entry_disabled_by",
|
|
"device_disabled_by",
|
|
"entity_disabled_by",
|
|
"setup_result",
|
|
"minor_version_after_migration",
|
|
"config_entry_disabled_by_after_migration",
|
|
"device_disabled_by_after_migration",
|
|
"entity_disabled_by_after_migration",
|
|
),
|
|
[
|
|
# Config entry not disabled, update device and entity disabled by config entry
|
|
(
|
|
None,
|
|
DeviceEntryDisabler.CONFIG_ENTRY,
|
|
RegistryEntryDisabler.CONFIG_ENTRY,
|
|
True,
|
|
3,
|
|
None,
|
|
DeviceEntryDisabler.USER,
|
|
RegistryEntryDisabler.DEVICE,
|
|
),
|
|
(
|
|
None,
|
|
DeviceEntryDisabler.USER,
|
|
RegistryEntryDisabler.DEVICE,
|
|
True,
|
|
3,
|
|
None,
|
|
DeviceEntryDisabler.USER,
|
|
RegistryEntryDisabler.DEVICE,
|
|
),
|
|
(
|
|
None,
|
|
DeviceEntryDisabler.USER,
|
|
RegistryEntryDisabler.USER,
|
|
True,
|
|
3,
|
|
None,
|
|
DeviceEntryDisabler.USER,
|
|
RegistryEntryDisabler.USER,
|
|
),
|
|
(
|
|
None,
|
|
None,
|
|
None,
|
|
True,
|
|
3,
|
|
None,
|
|
None,
|
|
None,
|
|
),
|
|
# Config entry disabled, migration does not run
|
|
(
|
|
ConfigEntryDisabler.USER,
|
|
DeviceEntryDisabler.CONFIG_ENTRY,
|
|
RegistryEntryDisabler.CONFIG_ENTRY,
|
|
False,
|
|
2,
|
|
ConfigEntryDisabler.USER,
|
|
DeviceEntryDisabler.CONFIG_ENTRY,
|
|
RegistryEntryDisabler.CONFIG_ENTRY,
|
|
),
|
|
(
|
|
ConfigEntryDisabler.USER,
|
|
DeviceEntryDisabler.USER,
|
|
RegistryEntryDisabler.DEVICE,
|
|
False,
|
|
2,
|
|
ConfigEntryDisabler.USER,
|
|
DeviceEntryDisabler.USER,
|
|
RegistryEntryDisabler.DEVICE,
|
|
),
|
|
(
|
|
ConfigEntryDisabler.USER,
|
|
DeviceEntryDisabler.USER,
|
|
RegistryEntryDisabler.USER,
|
|
False,
|
|
2,
|
|
ConfigEntryDisabler.USER,
|
|
DeviceEntryDisabler.USER,
|
|
RegistryEntryDisabler.USER,
|
|
),
|
|
(
|
|
ConfigEntryDisabler.USER,
|
|
None,
|
|
None,
|
|
False,
|
|
2,
|
|
ConfigEntryDisabler.USER,
|
|
None,
|
|
None,
|
|
),
|
|
],
|
|
)
|
|
async def test_migrate_entry_from_v3_2(
|
|
hass: HomeAssistant,
|
|
device_registry: dr.DeviceRegistry,
|
|
entity_registry: er.EntityRegistry,
|
|
config_entry_disabled_by: ConfigEntryDisabler | None,
|
|
device_disabled_by: DeviceEntryDisabler | None,
|
|
entity_disabled_by: RegistryEntryDisabler | None,
|
|
setup_result: bool,
|
|
minor_version_after_migration: int,
|
|
config_entry_disabled_by_after_migration: ConfigEntryDisabler | None,
|
|
device_disabled_by_after_migration: ConfigEntryDisabler | None,
|
|
entity_disabled_by_after_migration: RegistryEntryDisabler | None,
|
|
) -> None:
|
|
"""Test migration from version 3.2."""
|
|
# Create a v3.2 config entry with conversation subentries
|
|
conversation_subentry_id = "blabla"
|
|
mock_config_entry = MockConfigEntry(
|
|
domain=DOMAIN,
|
|
data={CONF_URL: "http://localhost:11434"},
|
|
disabled_by=config_entry_disabled_by,
|
|
version=3,
|
|
minor_version=2,
|
|
subentries_data=[
|
|
{
|
|
"data": V1_TEST_OPTIONS,
|
|
"subentry_id": conversation_subentry_id,
|
|
"subentry_type": "conversation",
|
|
"title": "Ollama",
|
|
"unique_id": None,
|
|
},
|
|
{
|
|
"data": {"model": "llama3.2:latest"},
|
|
"subentry_type": "ai_task_data",
|
|
"title": "Ollama AI Task",
|
|
"unique_id": None,
|
|
},
|
|
],
|
|
)
|
|
mock_config_entry.add_to_hass(hass)
|
|
|
|
conversation_device = device_registry.async_get_or_create(
|
|
config_entry_id=mock_config_entry.entry_id,
|
|
config_subentry_id=conversation_subentry_id,
|
|
disabled_by=device_disabled_by,
|
|
identifiers={(DOMAIN, mock_config_entry.entry_id)},
|
|
name=mock_config_entry.title,
|
|
manufacturer="Ollama",
|
|
model="Ollama",
|
|
entry_type=dr.DeviceEntryType.SERVICE,
|
|
)
|
|
conversation_entity = entity_registry.async_get_or_create(
|
|
"conversation",
|
|
DOMAIN,
|
|
mock_config_entry.entry_id,
|
|
config_entry=mock_config_entry,
|
|
config_subentry_id=conversation_subentry_id,
|
|
disabled_by=entity_disabled_by,
|
|
device_id=conversation_device.id,
|
|
suggested_object_id="ollama",
|
|
)
|
|
|
|
# Verify initial state
|
|
assert mock_config_entry.version == 3
|
|
assert mock_config_entry.minor_version == 2
|
|
assert len(mock_config_entry.subentries) == 2
|
|
assert mock_config_entry.disabled_by == config_entry_disabled_by
|
|
assert conversation_device.disabled_by == device_disabled_by
|
|
assert conversation_entity.disabled_by == entity_disabled_by
|
|
|
|
# Run setup to trigger migration
|
|
with patch(
|
|
"homeassistant.components.ollama.async_setup_entry",
|
|
return_value=True,
|
|
):
|
|
result = await hass.config_entries.async_setup(mock_config_entry.entry_id)
|
|
assert result is setup_result
|
|
await hass.async_block_till_done()
|
|
|
|
# Verify migration completed
|
|
entries = hass.config_entries.async_entries(DOMAIN)
|
|
assert len(entries) == 1
|
|
entry = entries[0]
|
|
|
|
# Check version and subversion were updated
|
|
assert entry.version == 3
|
|
assert entry.minor_version == minor_version_after_migration
|
|
|
|
# Check the disabled_by flag on config entry, device and entity are as expected
|
|
conversation_device = device_registry.async_get(conversation_device.id)
|
|
conversation_entity = entity_registry.async_get(conversation_entity.entity_id)
|
|
assert mock_config_entry.disabled_by == config_entry_disabled_by_after_migration
|
|
assert conversation_device.disabled_by == device_disabled_by_after_migration
|
|
assert conversation_entity.disabled_by == entity_disabled_by_after_migration
|