Correct ollama config entry migration (#147858)

This commit is contained in:
Erik Montnemery 2025-07-01 10:45:17 +02:00 committed by Franck Nijhof
parent 725269ecda
commit 8b2f4f0f86
No known key found for this signature in database
GPG Key ID: AB33ADACE7101952
3 changed files with 186 additions and 0 deletions

View File

@ -148,4 +148,34 @@ async def async_migrate_integration(hass: HomeAssistant) -> None:
title=DEFAULT_NAME,
options={},
version=2,
minor_version=2,
)
async def async_migrate_entry(hass: HomeAssistant, entry: OllamaConfigEntry) -> bool:
"""Migrate entry."""
_LOGGER.debug("Migrating from version %s:%s", entry.version, entry.minor_version)
if entry.version > 2:
# This means the user has downgraded from a future version
return False
if entry.version == 2 and entry.minor_version == 1:
# Correct broken device migration in Home Assistant Core 2025.7.0b0-2025.7.0b1
device_registry = dr.async_get(hass)
for device in dr.async_entries_for_config_entry(
device_registry, entry.entry_id
):
device_registry.async_update_device(
device.id,
remove_config_entry_id=entry.entry_id,
remove_config_subentry_id=None,
)
hass.config_entries.async_update_entry(entry, minor_version=2)
_LOGGER.debug(
"Migration to version %s:%s successful", entry.version, entry.minor_version
)
return True

View File

@ -73,6 +73,7 @@ class OllamaConfigFlow(ConfigFlow, domain=DOMAIN):
"""Handle a config flow for Ollama."""
VERSION = 2
MINOR_VERSION = 2
def __init__(self) -> None:
"""Initialize config flow."""

View File

@ -7,6 +7,7 @@ import pytest
from homeassistant.components import ollama
from homeassistant.components.ollama.const import DOMAIN
from homeassistant.config_entries import ConfigSubentryData
from homeassistant.core import HomeAssistant
from homeassistant.helpers import device_registry as dr, entity_registry as er
from homeassistant.setup import async_setup_component
@ -81,6 +82,7 @@ async def test_migration_from_v1_to_v2(
await hass.config_entries.async_setup(mock_config_entry.entry_id)
assert mock_config_entry.version == 2
assert mock_config_entry.minor_version == 2
assert mock_config_entry.data == TEST_USER_DATA
assert mock_config_entry.options == {}
@ -186,6 +188,7 @@ async def test_migration_from_v1_to_v2_with_multiple_urls(
for idx, entry in enumerate(entries):
assert entry.version == 2
assert entry.minor_version == 2
assert not entry.options
assert len(entry.subentries) == 1
subentry = list(entry.subentries.values())[0]
@ -273,6 +276,7 @@ async def test_migration_from_v1_to_v2_with_same_urls(
entry = entries[0]
assert entry.version == 2
assert entry.minor_version == 2
assert not entry.options
assert len(entry.subentries) == 2 # Two subentries from the two original entries
@ -295,3 +299,154 @@ async def test_migration_from_v1_to_v2_with_same_urls(
assert dev.config_entries_subentries == {
mock_config_entry.entry_id: {subentry.subentry_id}
}
async def test_migration_from_v2_1_to_v2_2(
hass: HomeAssistant,
device_registry: dr.DeviceRegistry,
entity_registry: er.EntityRegistry,
) -> None:
"""Test migration from version 2.1 to version 2.2.
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=TEST_USER_DATA,
entry_id="mock_entry_id",
version=2,
minor_version=1,
subentries_data=[
ConfigSubentryData(
data=TEST_OPTIONS,
subentry_id="mock_id_1",
subentry_type="conversation",
title="Ollama",
unique_id=None,
),
ConfigSubentryData(
data=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 == 2
assert entry.minor_version == 2
assert not entry.options
assert entry.title == "Ollama"
assert len(entry.subentries) == 2
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 == 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}
}