From b066877453ebe8d840f699721f0f644254e68acc Mon Sep 17 00:00:00 2001
From: Anders Melchiorsen <amelchio@nogoto.net>
Date: Sun, 18 Nov 2018 19:51:17 +0100
Subject: [PATCH] Allow unloading of LIFX config entry (#18535)

---
 homeassistant/components/lifx/__init__.py | 14 +++++-
 homeassistant/components/light/lifx.py    | 56 +++++++++++++++--------
 requirements_all.txt                      |  2 +-
 3 files changed, 50 insertions(+), 22 deletions(-)

diff --git a/homeassistant/components/lifx/__init__.py b/homeassistant/components/lifx/__init__.py
index 1ca6c00b23a..52df3d47ca1 100644
--- a/homeassistant/components/lifx/__init__.py
+++ b/homeassistant/components/lifx/__init__.py
@@ -8,7 +8,7 @@ from homeassistant.components.light import DOMAIN as LIGHT_DOMAIN
 
 
 DOMAIN = 'lifx'
-REQUIREMENTS = ['aiolifx==0.6.5']
+REQUIREMENTS = ['aiolifx==0.6.6']
 
 CONF_SERVER = 'server'
 CONF_BROADCAST = 'broadcast'
@@ -25,6 +25,8 @@ CONFIG_SCHEMA = vol.Schema({
     }
 }, extra=vol.ALLOW_EXTRA)
 
+DATA_LIFX_MANAGER = 'lifx_manager'
+
 
 async def async_setup(hass, config):
     """Set up the LIFX component."""
@@ -43,6 +45,16 @@ async def async_setup_entry(hass, entry):
     """Set up LIFX from a config entry."""
     hass.async_create_task(hass.config_entries.async_forward_entry_setup(
         entry, LIGHT_DOMAIN))
+
+    return True
+
+
+async def async_unload_entry(hass, entry):
+    """Unload a config entry."""
+    hass.data.pop(DATA_LIFX_MANAGER).cleanup()
+
+    await hass.config_entries.async_forward_entry_unload(entry, LIGHT_DOMAIN)
+
     return True
 
 
diff --git a/homeassistant/components/light/lifx.py b/homeassistant/components/light/lifx.py
index f346f88c42b..8951b2876a2 100644
--- a/homeassistant/components/light/lifx.py
+++ b/homeassistant/components/light/lifx.py
@@ -22,7 +22,7 @@ from homeassistant.components.light import (
     SUPPORT_TRANSITION, VALID_BRIGHTNESS, VALID_BRIGHTNESS_PCT, Light,
     preprocess_turn_on_alternatives)
 from homeassistant.components.lifx import (
-    DOMAIN as LIFX_DOMAIN, CONF_SERVER, CONF_BROADCAST)
+    DOMAIN as LIFX_DOMAIN, DATA_LIFX_MANAGER, CONF_SERVER, CONF_BROADCAST)
 from homeassistant.const import ATTR_ENTITY_ID, EVENT_HOMEASSISTANT_STOP
 from homeassistant.core import callback
 import homeassistant.helpers.config_validation as cv
@@ -155,27 +155,10 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
             interfaces = [{}]
 
     lifx_manager = LIFXManager(hass, async_add_entities)
+    hass.data[DATA_LIFX_MANAGER] = lifx_manager
 
     for interface in interfaces:
-        kwargs = {'discovery_interval': DISCOVERY_INTERVAL}
-        broadcast_ip = interface.get(CONF_BROADCAST)
-        if broadcast_ip:
-            kwargs['broadcast_ip'] = broadcast_ip
-        lifx_discovery = aiolifx().LifxDiscovery(
-            hass.loop, lifx_manager, **kwargs)
-
-        kwargs = {}
-        listen_ip = interface.get(CONF_SERVER)
-        if listen_ip:
-            kwargs['listen_ip'] = listen_ip
-        lifx_discovery.start(**kwargs)
-
-    @callback
-    def cleanup(event):
-        """Clean up resources."""
-        lifx_discovery.cleanup()
-
-    hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, cleanup)
+        lifx_manager.start_discovery(interface)
 
     return True
 
@@ -226,10 +209,43 @@ class LIFXManager:
         self.hass = hass
         self.async_add_entities = async_add_entities
         self.effects_conductor = aiolifx_effects().Conductor(loop=hass.loop)
+        self.discoveries = []
+        self.cleanup_unsub = self.hass.bus.async_listen(
+            EVENT_HOMEASSISTANT_STOP,
+            self.cleanup)
 
         self.register_set_state()
         self.register_effects()
 
+    def start_discovery(self, interface):
+        """Start discovery on a network interface."""
+        kwargs = {'discovery_interval': DISCOVERY_INTERVAL}
+        broadcast_ip = interface.get(CONF_BROADCAST)
+        if broadcast_ip:
+            kwargs['broadcast_ip'] = broadcast_ip
+        lifx_discovery = aiolifx().LifxDiscovery(
+            self.hass.loop, self, **kwargs)
+
+        kwargs = {}
+        listen_ip = interface.get(CONF_SERVER)
+        if listen_ip:
+            kwargs['listen_ip'] = listen_ip
+        lifx_discovery.start(**kwargs)
+
+        self.discoveries.append(lifx_discovery)
+
+    @callback
+    def cleanup(self, event=None):
+        """Release resources."""
+        self.cleanup_unsub()
+
+        for discovery in self.discoveries:
+            discovery.cleanup()
+
+        for service in [SERVICE_LIFX_SET_STATE, SERVICE_EFFECT_STOP,
+                        SERVICE_EFFECT_PULSE, SERVICE_EFFECT_COLORLOOP]:
+            self.hass.services.async_remove(DOMAIN, service)
+
     def register_set_state(self):
         """Register the LIFX set_state service call."""
         async def service_handler(service):
diff --git a/requirements_all.txt b/requirements_all.txt
index 76cb14fdf1a..67b0bebdd97 100644
--- a/requirements_all.txt
+++ b/requirements_all.txt
@@ -111,7 +111,7 @@ aiohue==1.5.0
 aioimaplib==0.7.13
 
 # homeassistant.components.lifx
-aiolifx==0.6.5
+aiolifx==0.6.6
 
 # homeassistant.components.light.lifx
 aiolifx_effects==0.2.1