Allow platform unloading (#13784)

* Allow platform unloading

* Add tests

* Add last test

homeassistant/components/hue/__init__.py

@@ -131,3 +131,9 @@ async def async_setup_entry(hass, entry):
     bridge = HueBridge(hass, entry, allow_unreachable, allow_groups)
     hass.data[DOMAIN][host] = bridge
     return await bridge.async_setup()
+
+
+async def async_unload_entry(hass, entry):
+    """Unload a config entry."""
+    bridge = hass.data[DOMAIN].pop(entry.data['host'])
+    return await bridge.async_reset()

homeassistant/components/hue/bridge.py

@@ -30,6 +30,7 @@ class HueBridge(object):
         self.allow_groups = allow_groups
         self.available = True
         self.api = None
+        self._cancel_retry_setup = None
 
     @property
     def host(self):
@@ -67,8 +68,8 @@ class HueBridge(object):
                     # This feels hacky, we should find a better way to do this
                     self.config_entry.state = config_entries.ENTRY_STATE_LOADED
 
-            # Unhandled edge case: cancel this if we discover bridge on new IP
-            hass.helpers.event.async_call_later(retry_delay, retry_setup)
+            self._cancel_retry_setup = hass.helpers.event.async_call_later(
+                retry_delay, retry_setup)
 
             return False
 
@@ -77,7 +78,7 @@ class HueBridge(object):
                              host)
             return False
 
-        hass.async_add_job(hass.config_entries.async_forward_entry(
+        hass.async_add_job(hass.config_entries.async_forward_entry_setup(
             self.config_entry, 'light'))
 
         hass.services.async_register(
@@ -86,6 +87,34 @@ class HueBridge(object):
 
         return True
 
+    async def async_reset(self):
+        """Reset this bridge to default state.
+
+        Will cancel any scheduled setup retry and will unload
+        the config entry.
+        """
+        # The bridge can be in 3 states:
+        #  - Setup was successful, self.api is not None
+        #  - Authentication was wrong, self.api is None, not retrying setup.
+        #  - Host was down. self.api is None, we're retrying setup
+
+        # If we have a retry scheduled, we were never setup.
+        if self._cancel_retry_setup is not None:
+            self._cancel_retry_setup()
+            self._cancel_retry_setup = None
+            return True
+
+        # If the authentication was wrong.
+        if self.api is None:
+            return True
+
+        self.hass.services.async_remove(DOMAIN, SERVICE_HUE_SCENE)
+
+        # If setup was successful, we set api variable, forwarded entry and
+        # register service
+        return await self.hass.config_entries.async_forward_entry_unload(
+            self.config_entry, 'light')
+
     async def hue_activate_scene(self, call, updated=False):
         """Service to call directly into bridge to set scenes."""
         group_name = call.data[ATTR_GROUP_NAME]

homeassistant/components/light/__init__.py

@@ -393,6 +393,11 @@ async def async_setup_entry(hass, entry):
     return await hass.data[DOMAIN].async_setup_entry(entry)
 
 
+async def async_unload_entry(hass, entry):
+    """Unload a config entry."""
+    return await hass.data[DOMAIN].async_unload_entry(entry)
+
+
 class Profiles:
     """Representation of available color profiles."""
 

homeassistant/config_entries.py

@@ -203,12 +203,13 @@ class ConfigEntry:
         else:
             self.state = ENTRY_STATE_SETUP_ERROR
 
-    async def async_unload(self, hass):
+    async def async_unload(self, hass, *, component=None):
         """Unload an entry.
 
         Returns if unload is possible and was successful.
         """
-        component = getattr(hass.components, self.domain)
+        if component is None:
+            component = getattr(hass.components, self.domain)
 
         supports_unload = hasattr(component, 'async_unload_entry')
 
@@ -220,13 +221,13 @@ class ConfigEntry:
 
             if not isinstance(result, bool):
                 _LOGGER.error('%s.async_unload_entry did not return boolean',
-                              self.domain)
+                              component.DOMAIN)
                 result = False
 
             return result
         except Exception:  # pylint: disable=broad-except
             _LOGGER.exception('Error unloading entry %s for %s',
-                              self.title, self.domain)
+                              self.title, component.DOMAIN)
             self.state = ENTRY_STATE_FAILED_UNLOAD
             return False
 
@@ -326,7 +327,7 @@ class ConfigEntries:
         entries = await self.hass.async_add_job(load_json, path)
         self._entries = [ConfigEntry(**entry) for entry in entries]
 
-    async def async_forward_entry(self, entry, component):
+    async def async_forward_entry_setup(self, entry, component):
         """Forward the setup of an entry to a different component.
 
         By default an entry is setup with the component it belongs to. If that
@@ -347,6 +348,15 @@ class ConfigEntries:
         await entry.async_setup(
             self.hass, component=getattr(self.hass.components, component))
 
+    async def async_forward_entry_unload(self, entry, component):
+        """Forward the unloading of an entry to a different component."""
+        # It was never loaded.
+        if component not in self.hass.config.components:
+            return True
+
+        await entry.async_unload(
+            self.hass, component=getattr(self.hass.components, component))
+
     async def _async_add_entry(self, entry):
         """Add an entry."""
         self._entries.append(entry)

homeassistant/helpers/entity_component.py

@@ -113,6 +113,18 @@ class EntityComponent(object):
 
         return await self._platforms[key].async_setup_entry(config_entry)
 
+    async def async_unload_entry(self, config_entry):
+        """Unload a config entry."""
+        key = config_entry.entry_id
+
+        platform = self._platforms.pop(key, None)
+
+        if platform is None:
+            raise ValueError('Config entry was never loaded!')
+
+        await platform.async_reset()
+        return True
+
     @callback
     def async_extract_from_service(self, service, expand_group=True):
         """Extract all known and available entities from a service call.

homeassistant/helpers/entity_platform.py

@@ -43,7 +43,10 @@ class EntityPlatform(object):
         self.config_entry = None
         self.entities = {}
         self._tasks = []
+        # Method to cancel the state change listener
         self._async_unsub_polling = None
+        # Method to cancel the retry of setup
+        self._async_cancel_retry_setup = None
         self._process_updates = asyncio.Lock(loop=hass.loop)
 
         # Platform is None for the EntityComponent "catch-all" EntityPlatform
@@ -145,10 +148,12 @@ class EntityPlatform(object):
 
             async def setup_again(now):
                 """Run setup again."""
+                self._async_cancel_retry_setup = None
                 await self._async_setup_platform(
                     async_create_setup_task, tries)
 
-            async_call_later(hass, wait_time, setup_again)
+            self._async_cancel_retry_setup = \
+                async_call_later(hass, wait_time, setup_again)
             return False
         except asyncio.TimeoutError:
             logger.error(
@@ -315,6 +320,10 @@ class EntityPlatform(object):
 
         This method must be run in the event loop.
         """
+        if self._async_cancel_retry_setup is not None:
+            self._async_cancel_retry_setup()
+            self._async_cancel_retry_setup = None
+
         if not self.entities:
             return
 

tests/components/hue/test_bridge.py

@@ -18,8 +18,8 @@ async def test_bridge_setup():
         assert await hue_bridge.async_setup() is True
 
     assert hue_bridge.api is api
-    assert len(hass.config_entries.async_forward_entry.mock_calls) == 1
-    assert hass.config_entries.async_forward_entry.mock_calls[0][1] == \
+    assert len(hass.config_entries.async_forward_entry_setup.mock_calls) == 1
+    assert hass.config_entries.async_forward_entry_setup.mock_calls[0][1] == \
         (entry, 'light')
 
 
@@ -54,3 +54,60 @@ async def test_bridge_setup_timeout(hass):
     assert len(hass.helpers.event.async_call_later.mock_calls) == 1
     # Assert we are going to wait 2 seconds
     assert hass.helpers.event.async_call_later.mock_calls[0][1][0] == 2
+
+
+async def test_reset_cancels_retry_setup():
+    """Test resetting a bridge while we're waiting to retry setup."""
+    hass = Mock()
+    entry = Mock()
+    entry.data = {'host': '1.2.3.4', 'username': 'mock-username'}
+    hue_bridge = bridge.HueBridge(hass, entry, False, False)
+
+    with patch.object(bridge, 'get_bridge', side_effect=errors.CannotConnect):
+        assert await hue_bridge.async_setup() is False
+
+    mock_call_later = hass.helpers.event.async_call_later
+
+    assert len(mock_call_later.mock_calls) == 1
+
+    assert await hue_bridge.async_reset()
+
+    assert len(mock_call_later.mock_calls) == 2
+    assert len(mock_call_later.return_value.mock_calls) == 1
+
+
+async def test_reset_if_entry_had_wrong_auth():
+    """Test calling reset when the entry contained wrong auth."""
+    hass = Mock()
+    entry = Mock()
+    entry.data = {'host': '1.2.3.4', 'username': 'mock-username'}
+    hue_bridge = bridge.HueBridge(hass, entry, False, False)
+
+    with patch.object(bridge, 'get_bridge',
+                      side_effect=errors.AuthenticationRequired):
+        assert await hue_bridge.async_setup() is False
+
+    assert len(hass.async_add_job.mock_calls) == 1
+
+    assert await hue_bridge.async_reset()
+
+
+async def test_reset_unloads_entry_if_setup():
+    """Test calling reset while the entry has been setup."""
+    hass = Mock()
+    entry = Mock()
+    entry.data = {'host': '1.2.3.4', 'username': 'mock-username'}
+    hue_bridge = bridge.HueBridge(hass, entry, False, False)
+
+    with patch.object(bridge, 'get_bridge', return_value=mock_coro(Mock())):
+        assert await hue_bridge.async_setup() is True
+
+    assert len(hass.services.async_register.mock_calls) == 1
+    assert len(hass.config_entries.async_forward_entry_setup.mock_calls) == 1
+
+    hass.config_entries.async_forward_entry_unload.return_value = \
+        mock_coro(True)
+    assert await hue_bridge.async_reset()
+
+    assert len(hass.config_entries.async_forward_entry_unload.mock_calls) == 1
+    assert len(hass.services.async_remove.mock_calls) == 1

tests/components/hue/test_init.py

@@ -167,3 +167,22 @@ async def test_config_passed_to_config_entry(hass):
     assert p_entry is entry
     assert p_allow_unreachable is True
     assert p_allow_groups is False
+
+
+async def test_unload_entry(hass):
+    """Test being able to unload an entry."""
+    entry = MockConfigEntry(domain=hue.DOMAIN, data={
+        'host': '0.0.0.0',
+    })
+    entry.add_to_hass(hass)
+
+    with patch.object(hue, 'HueBridge') as mock_bridge:
+        mock_bridge.return_value.async_setup.return_value = mock_coro(True)
+        assert await async_setup_component(hass, hue.DOMAIN, {}) is True
+
+    assert len(mock_bridge.return_value.mock_calls) == 1
+
+    mock_bridge.return_value.async_reset.return_value = mock_coro(True)
+    assert await hue.async_unload_entry(hass, entry)
+    assert len(mock_bridge.return_value.async_reset.mock_calls) == 1
+    assert hass.data[hue.DOMAIN] == {}

tests/components/light/test_hue.py

@@ -200,7 +200,7 @@ async def setup_bridge(hass, mock_bridge):
     config_entry = config_entries.ConfigEntry(1, hue.DOMAIN, 'Mock Title', {
         'host': 'mock-host'
     }, 'test')
-    await hass.config_entries.async_forward_entry(config_entry, 'light')
+    await hass.config_entries.async_forward_entry_setup(config_entry, 'light')
     # To flush out the service call to update the group
     await hass.async_block_till_done()
 

tests/helpers/test_entity_component.py

@@ -376,3 +376,34 @@ async def test_setup_entry_fails_duplicate(hass):
 
     with pytest.raises(ValueError):
         await component.async_setup_entry(entry)
+
+
+async def test_unload_entry_resets_platform(hass):
+    """Test unloading an entry removes all entities."""
+    mock_setup_entry = Mock(return_value=mock_coro(True))
+    loader.set_component(
+        'test_domain.entry_domain',
+        MockPlatform(async_setup_entry=mock_setup_entry))
+
+    component = EntityComponent(_LOGGER, DOMAIN, hass)
+    entry = MockConfigEntry(domain='entry_domain')
+
+    assert await component.async_setup_entry(entry)
+    assert len(mock_setup_entry.mock_calls) == 1
+    add_entities = mock_setup_entry.mock_calls[0][1][2]
+    add_entities([MockEntity()])
+    await hass.async_block_till_done()
+
+    assert len(hass.states.async_entity_ids()) == 1
+
+    assert await component.async_unload_entry(entry)
+    assert len(hass.states.async_entity_ids()) == 0
+
+
+async def test_unload_entry_fails_if_never_loaded(hass):
+    """."""
+    component = EntityComponent(_LOGGER, DOMAIN, hass)
+    entry = MockConfigEntry(domain='entry_domain')
+
+    with pytest.raises(ValueError):
+        await component.async_unload_entry(entry)

tests/helpers/test_entity_platform.py

@@ -555,3 +555,29 @@ async def test_setup_entry_platform_not_ready(hass, caplog):
     assert len(async_setup_entry.mock_calls) == 1
     assert 'Platform test not ready yet' in caplog.text
     assert len(mock_call_later.mock_calls) == 1
+
+
+async def test_reset_cancels_retry_setup(hass):
+    """Test that resetting a platform will cancel scheduled a setup retry."""
+    async_setup_entry = Mock(side_effect=PlatformNotReady)
+    platform = MockPlatform(
+        async_setup_entry=async_setup_entry
+    )
+    config_entry = MockConfigEntry()
+    ent_platform = MockEntityPlatform(
+        hass,
+        platform_name=config_entry.domain,
+        platform=platform
+    )
+
+    with patch.object(entity_platform, 'async_call_later') as mock_call_later:
+        assert not await ent_platform.async_setup_entry(config_entry)
+
+    assert len(mock_call_later.mock_calls) == 1
+    assert len(mock_call_later.return_value.mock_calls) == 0
+    assert ent_platform._async_cancel_retry_setup is not None
+
+    await ent_platform.async_reset()
+
+    assert len(mock_call_later.return_value.mock_calls) == 1
+    assert ent_platform._async_cancel_retry_setup is None

tests/test_config_entries.py

@@ -405,7 +405,7 @@ async def test_forward_entry_sets_up_component(hass):
         'forwarded',
         MockModule('forwarded', async_setup_entry=mock_forwarded_setup_entry))
 
-    await hass.config_entries.async_forward_entry(entry, 'forwarded')
+    await hass.config_entries.async_forward_entry_setup(entry, 'forwarded')
     assert len(mock_original_setup_entry.mock_calls) == 0
     assert len(mock_forwarded_setup_entry.mock_calls) == 1
 
@@ -422,6 +422,6 @@ async def test_forward_entry_does_not_setup_entry_if_setup_fails(hass):
         async_setup_entry=mock_setup_entry,
     ))
 
-    await hass.config_entries.async_forward_entry(entry, 'forwarded')
+    await hass.config_entries.async_forward_entry_setup(entry, 'forwarded')
     assert len(mock_setup.mock_calls) == 1
     assert len(mock_setup_entry.mock_calls) == 0