Config Entry migrations (#20888)

* Updated per review feedback.

* Fixed line length

* Review comments and lint error

* Fixed mypy typeing error

* Moved migration logic to setup

* Use new migration error state

* Fix bug and ignore mypy type error

* Removed SmartThings example and added unit tests.

* Fixed test comments.

homeassistant/config_entries.py

@@ -7,7 +7,11 @@ component.
 During startup, Home Assistant will setup the entries during the normal setup
 of a component. It will first call the normal setup and then call the method
 `async_setup_entry(hass, entry)` for each entry. The same method is called when
-Home Assistant is running while a config entry is created.
+Home Assistant is running while a config entry is created.  If the version of
+the config entry does not match that of the flow handler, setup will
+call the method `async_migrate_entry(hass, entry)` with the expectation that
+the entry be brought to the current version.  Return `True` to indicate
+migration was successful, otherwise `False`.
 
 ## Config Flows
 
@@ -116,6 +120,7 @@ If the result of the step is to show a form, the user will be able to continue
 the flow from the config panel.
 """
 import logging
+import functools
 import uuid
 from typing import Set, Optional, List, Dict  # noqa pylint: disable=unused-import
 
@@ -188,6 +193,8 @@ SAVE_DELAY = 1
 ENTRY_STATE_LOADED = 'loaded'
 # There was an error while trying to set up this config entry
 ENTRY_STATE_SETUP_ERROR = 'setup_error'
+# There was an error while trying to migrate the config entry to a new version
+ENTRY_STATE_MIGRATION_ERROR = 'migration_error'
 # The config entry was not ready to be set up yet, but might be later
 ENTRY_STATE_SETUP_RETRY = 'setup_retry'
 # The config entry has not been loaded
@@ -256,6 +263,12 @@ class ConfigEntry:
         if component is None:
             component = getattr(hass.components, self.domain)
 
+        # Perform migration
+        if component.DOMAIN == self.domain:
+            if not await self.async_migrate(hass):
+                self.state = ENTRY_STATE_MIGRATION_ERROR
+                return
+
         try:
             result = await component.async_setup_entry(hass, self)
 
@@ -332,6 +345,45 @@ class ConfigEntry:
                 self.state = ENTRY_STATE_FAILED_UNLOAD
             return False
 
+    async def async_migrate(self, hass: HomeAssistant) -> bool:
+        """Migrate an entry.
+
+        Returns True if config entry is up-to-date or has been migrated.
+        """
+        handler = HANDLERS.get(self.domain)
+        if handler is None:
+            _LOGGER.error("Flow handler not found for entry %s for %s",
+                          self.title, self.domain)
+            return False
+        # Handler may be a partial
+        while isinstance(handler, functools.partial):
+            handler = handler.func
+
+        if self.version == handler.VERSION:
+            return True
+
+        component = getattr(hass.components, self.domain)
+        supports_migrate = hasattr(component, 'async_migrate_entry')
+        if not supports_migrate:
+            _LOGGER.error("Migration handler not found for entry %s for %s",
+                          self.title, self.domain)
+            return False
+
+        try:
+            result = await component.async_migrate_entry(hass, self)
+            if not isinstance(result, bool):
+                _LOGGER.error('%s.async_migrate_entry did not return boolean',
+                              self.domain)
+                return False
+            if result:
+                # pylint: disable=protected-access
+                hass.config_entries._async_schedule_save()  # type: ignore
+            return result
+        except Exception:  # pylint: disable=broad-except
+            _LOGGER.exception('Error migrating entry %s for %s',
+                              self.title, component.DOMAIN)
+            return False
+
     def as_dict(self):
         """Return dictionary version of this entry."""
         return {

tests/common.py

@@ -451,7 +451,8 @@ class MockModule:
     def __init__(self, domain=None, dependencies=None, setup=None,
                  requirements=None, config_schema=None, platform_schema=None,
                  platform_schema_base=None, async_setup=None,
-                 async_setup_entry=None, async_unload_entry=None):
+                 async_setup_entry=None, async_unload_entry=None,
+                 async_migrate_entry=None):
         """Initialize the mock module."""
         self.DOMAIN = domain
         self.DEPENDENCIES = dependencies or []
@@ -482,6 +483,9 @@ class MockModule:
         if async_unload_entry is not None:
             self.async_unload_entry = async_unload_entry
 
+        if async_migrate_entry is not None:
+            self.async_migrate_entry = async_migrate_entry
+
 
 class MockPlatform:
     """Provide a fake platform."""
@@ -602,7 +606,7 @@ class MockToggleDevice(entity.ToggleEntity):
 class MockConfigEntry(config_entries.ConfigEntry):
     """Helper for creating config entries that adds some defaults."""
 
-    def __init__(self, *, domain='test', data=None, version=0, entry_id=None,
+    def __init__(self, *, domain='test', data=None, version=1, entry_id=None,
                  source=config_entries.SOURCE_USER, title='Mock Title',
                  state=None,
                  connection_class=config_entries.CONN_CLASS_UNKNOWN):

tests/test_config_entries.py

@@ -15,6 +15,14 @@ from tests.common import (
     MockPlatform, MockEntity)
 
 
+@config_entries.HANDLERS.register('test')
+@config_entries.HANDLERS.register('comp')
+class MockFlowHandler(config_entries.ConfigFlow):
+    """Define a mock flow handler."""
+
+    VERSION = 1
+
+
 @pytest.fixture
 def manager(hass):
     """Fixture of a loaded config manager."""
@@ -25,10 +33,117 @@ def manager(hass):
     return manager
 
 
-@asyncio.coroutine
+async def test_call_setup_entry(hass):
-def test_call_setup_entry(hass):
     """Test we call <component>.setup_entry."""
-    MockConfigEntry(domain='comp').add_to_hass(hass)
+    entry = MockConfigEntry(domain='comp')
+    entry.add_to_hass(hass)
+
+    mock_setup_entry = MagicMock(return_value=mock_coro(True))
+    mock_migrate_entry = MagicMock(return_value=mock_coro(True))
+
+    loader.set_component(
+        hass, 'comp',
+        MockModule('comp', async_setup_entry=mock_setup_entry,
+                   async_migrate_entry=mock_migrate_entry))
+
+    result = await async_setup_component(hass, 'comp', {})
+    assert result
+    assert len(mock_migrate_entry.mock_calls) == 0
+    assert len(mock_setup_entry.mock_calls) == 1
+    assert entry.state == config_entries.ENTRY_STATE_LOADED
+
+
+async def test_call_async_migrate_entry(hass):
+    """Test we call <component>.async_migrate_entry when version mismatch."""
+    entry = MockConfigEntry(domain='comp')
+    entry.version = 2
+    entry.add_to_hass(hass)
+
+    mock_migrate_entry = MagicMock(return_value=mock_coro(True))
+    mock_setup_entry = MagicMock(return_value=mock_coro(True))
+
+    loader.set_component(
+        hass, 'comp',
+        MockModule('comp', async_setup_entry=mock_setup_entry,
+                   async_migrate_entry=mock_migrate_entry))
+
+    result = await async_setup_component(hass, 'comp', {})
+    assert result
+    assert len(mock_migrate_entry.mock_calls) == 1
+    assert len(mock_setup_entry.mock_calls) == 1
+    assert entry.state == config_entries.ENTRY_STATE_LOADED
+
+
+async def test_call_async_migrate_entry_failure_false(hass):
+    """Test migration fails if returns false."""
+    entry = MockConfigEntry(domain='comp')
+    entry.version = 2
+    entry.add_to_hass(hass)
+
+    mock_migrate_entry = MagicMock(return_value=mock_coro(False))
+    mock_setup_entry = MagicMock(return_value=mock_coro(True))
+
+    loader.set_component(
+        hass, 'comp',
+        MockModule('comp', async_setup_entry=mock_setup_entry,
+                   async_migrate_entry=mock_migrate_entry))
+
+    result = await async_setup_component(hass, 'comp', {})
+    assert result
+    assert len(mock_migrate_entry.mock_calls) == 1
+    assert len(mock_setup_entry.mock_calls) == 0
+    assert entry.state == config_entries.ENTRY_STATE_MIGRATION_ERROR
+
+
+async def test_call_async_migrate_entry_failure_exception(hass):
+    """Test migration fails if exception raised."""
+    entry = MockConfigEntry(domain='comp')
+    entry.version = 2
+    entry.add_to_hass(hass)
+
+    mock_migrate_entry = MagicMock(
+        return_value=mock_coro(exception=Exception))
+    mock_setup_entry = MagicMock(return_value=mock_coro(True))
+
+    loader.set_component(
+        hass, 'comp',
+        MockModule('comp', async_setup_entry=mock_setup_entry,
+                   async_migrate_entry=mock_migrate_entry))
+
+    result = await async_setup_component(hass, 'comp', {})
+    assert result
+    assert len(mock_migrate_entry.mock_calls) == 1
+    assert len(mock_setup_entry.mock_calls) == 0
+    assert entry.state == config_entries.ENTRY_STATE_MIGRATION_ERROR
+
+
+async def test_call_async_migrate_entry_failure_not_bool(hass):
+    """Test migration fails if boolean not returned."""
+    entry = MockConfigEntry(domain='comp')
+    entry.version = 2
+    entry.add_to_hass(hass)
+
+    mock_migrate_entry = MagicMock(
+        return_value=mock_coro())
+    mock_setup_entry = MagicMock(return_value=mock_coro(True))
+
+    loader.set_component(
+        hass, 'comp',
+        MockModule('comp', async_setup_entry=mock_setup_entry,
+                   async_migrate_entry=mock_migrate_entry))
+
+    result = await async_setup_component(hass, 'comp', {})
+    assert result
+    assert len(mock_migrate_entry.mock_calls) == 1
+    assert len(mock_setup_entry.mock_calls) == 0
+    assert entry.state == config_entries.ENTRY_STATE_MIGRATION_ERROR
+
+
+async def test_call_async_migrate_entry_failure_not_supported(hass):
+    """Test migration fails if async_migrate_entry not implemented."""
+    entry = MockConfigEntry(domain='comp')
+    entry.version = 2
+    entry.add_to_hass(hass)
 
     mock_setup_entry = MagicMock(return_value=mock_coro(True))
 
@@ -36,9 +151,10 @@ def test_call_setup_entry(hass):
         hass, 'comp',
         MockModule('comp', async_setup_entry=mock_setup_entry))
 
-    result = yield from async_setup_component(hass, 'comp', {})
+    result = await async_setup_component(hass, 'comp', {})
     assert result
-    assert len(mock_setup_entry.mock_calls) == 1
+    assert len(mock_setup_entry.mock_calls) == 0
+    assert entry.state == config_entries.ENTRY_STATE_MIGRATION_ERROR
 
 
 async def test_remove_entry(hass, manager):