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.

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
-def test_call_setup_entry(hass):
+async 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):