Fix unloading an entry can leave states around (#17786)

* Add test that tests unloading on remove

* Add more test things

* Untangle entity remove code from entity platform

* Don't add default implementation of async_will_remove

* Keep entity weakref alive

tests/test_config_entries.py

@@ -11,7 +11,8 @@ from homeassistant.setup import async_setup_component
 from homeassistant.util import dt
 
 from tests.common import (
-    MockModule, mock_coro, MockConfigEntry, async_fire_time_changed)
+    MockModule, mock_coro, MockConfigEntry, async_fire_time_changed,
+    MockPlatform, MockEntity)
 
 
 @pytest.fixture
@@ -40,35 +41,87 @@ def test_call_setup_entry(hass):
     assert len(mock_setup_entry.mock_calls) == 1
 
 
-@asyncio.coroutine
-def test_remove_entry(hass, manager):
+async def test_remove_entry(hass, manager):
     """Test that we can remove an entry."""
-    mock_unload_entry = MagicMock(return_value=mock_coro(True))
+    async def mock_setup_entry(hass, entry):
+        """Mock setting up entry."""
+        hass.loop.create_task(hass.config_entries.async_forward_entry_setup(
+            entry, 'light'))
+        return True
 
+    async def mock_unload_entry(hass, entry):
+        """Mock unloading an entry."""
+        result = await hass.config_entries.async_forward_entry_unload(
+            entry, 'light')
+        assert result
+        return result
+
+    entity = MockEntity(
+        unique_id='1234',
+        name='Test Entity',
+    )
+
+    async def mock_setup_entry_platform(hass, entry, async_add_entities):
+        """Mock setting up platform."""
+        async_add_entities([entity])
+
+    loader.set_component(hass, 'test', MockModule(
+        'test',
+        async_setup_entry=mock_setup_entry,
+        async_unload_entry=mock_unload_entry
+    ))
     loader.set_component(
-        hass, 'test',
-        MockModule('comp', async_unload_entry=mock_unload_entry))
+        hass, 'light.test',
+        MockPlatform(async_setup_entry=mock_setup_entry_platform))
 
     MockConfigEntry(domain='test', entry_id='test1').add_to_manager(manager)
-    MockConfigEntry(
+    entry = MockConfigEntry(
         domain='test',
         entry_id='test2',
-        state=config_entries.ENTRY_STATE_LOADED
-    ).add_to_manager(manager)
+    )
+    entry.add_to_manager(manager)
     MockConfigEntry(domain='test', entry_id='test3').add_to_manager(manager)
 
+    # Check all config entries exist
     assert [item.entry_id for item in manager.async_entries()] == \
         ['test1', 'test2', 'test3']
 
-    result = yield from manager.async_remove('test2')
+    # Setup entry
+    await entry.async_setup(hass)
+    await hass.async_block_till_done()
 
+    # Check entity state got added
+    assert hass.states.get('light.test_entity') is not None
+    # Group all_lights, light.test_entity
+    assert len(hass.states.async_all()) == 2
+
+    # Check entity got added to entity registry
+    ent_reg = await hass.helpers.entity_registry.async_get_registry()
+    assert len(ent_reg.entities) == 1
+    entity_entry = list(ent_reg.entities.values())[0]
+    assert entity_entry.config_entry_id == entry.entry_id
+
+    # Remove entry
+    result = await manager.async_remove('test2')
+    await hass.async_block_till_done()
+
+    # Check that unload went well and so no need to restart
     assert result == {
         'require_restart': False
     }
+
+    # Check that config entry was removed.
     assert [item.entry_id for item in manager.async_entries()] == \
         ['test1', 'test3']
 
-    assert len(mock_unload_entry.mock_calls) == 1
+    # Check that entity state has been removed
+    assert hass.states.get('light.test_entity') is None
+    # Just Group all_lights
+    assert len(hass.states.async_all()) == 1
+
+    # Check that entity registry entry no longer references config_entry_id
+    entity_entry = list(ent_reg.entities.values())[0]
+    assert entity_entry.config_entry_id is None
 
 
 @asyncio.coroutine