Explictly close the TPLink SmartDevice protocol on unload (#56743)

* Explictly close the TPLink SmartDevice protocol on unload

- There is a destructor that will eventually do this when
  the object gets gc. Its better to explictly do it at
  unload.

* fix coro mock
This commit is contained in:
J. Nick Koston 2021-09-28 11:36:45 -05:00 committed by GitHub
parent c95f7a5ba6
commit bc59387437
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 14 additions and 0 deletions

View File

@ -151,8 +151,10 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
hass_data: dict[str, Any] = hass.data[DOMAIN] hass_data: dict[str, Any] = hass.data[DOMAIN]
if entry.entry_id not in hass_data: if entry.entry_id not in hass_data:
return True return True
device: SmartDevice = hass.data[DOMAIN][entry.entry_id].device
if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS): if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
hass_data.pop(entry.entry_id) hass_data.pop(entry.entry_id)
await device.protocol.close()
return unload_ok return unload_ok

View File

@ -4,6 +4,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
from kasa import SmartBulb, SmartPlug, SmartStrip from kasa import SmartBulb, SmartPlug, SmartStrip
from kasa.exceptions import SmartDeviceException from kasa.exceptions import SmartDeviceException
from kasa.protocol import TPLinkSmartHomeProtocol
MODULE = "homeassistant.components.tplink" MODULE = "homeassistant.components.tplink"
MODULE_CONFIG_FLOW = "homeassistant.components.tplink.config_flow" MODULE_CONFIG_FLOW = "homeassistant.components.tplink.config_flow"
@ -14,6 +15,12 @@ MAC_ADDRESS = "aa:bb:cc:dd:ee:ff"
DEFAULT_ENTRY_TITLE = f"{ALIAS} {MODEL}" DEFAULT_ENTRY_TITLE = f"{ALIAS} {MODEL}"
def _mock_protocol() -> TPLinkSmartHomeProtocol:
protocol = MagicMock(auto_spec=TPLinkSmartHomeProtocol)
protocol.close = AsyncMock()
return protocol
def _mocked_bulb() -> SmartBulb: def _mocked_bulb() -> SmartBulb:
bulb = MagicMock(auto_spec=SmartBulb) bulb = MagicMock(auto_spec=SmartBulb)
bulb.update = AsyncMock() bulb.update = AsyncMock()
@ -36,6 +43,7 @@ def _mocked_bulb() -> SmartBulb:
bulb.set_brightness = AsyncMock() bulb.set_brightness = AsyncMock()
bulb.set_hsv = AsyncMock() bulb.set_hsv = AsyncMock()
bulb.set_color_temp = AsyncMock() bulb.set_color_temp = AsyncMock()
bulb.protocol = _mock_protocol()
return bulb return bulb
@ -55,6 +63,7 @@ def _mocked_plug() -> SmartPlug:
plug.hw_info = {"sw_ver": "1.0.0"} plug.hw_info = {"sw_ver": "1.0.0"}
plug.turn_off = AsyncMock() plug.turn_off = AsyncMock()
plug.turn_on = AsyncMock() plug.turn_on = AsyncMock()
plug.protocol = _mock_protocol()
return plug return plug
@ -74,14 +83,17 @@ def _mocked_strip() -> SmartStrip:
strip.hw_info = {"sw_ver": "1.0.0"} strip.hw_info = {"sw_ver": "1.0.0"}
strip.turn_off = AsyncMock() strip.turn_off = AsyncMock()
strip.turn_on = AsyncMock() strip.turn_on = AsyncMock()
strip.protocol = _mock_protocol()
plug0 = _mocked_plug() plug0 = _mocked_plug()
plug0.alias = "Plug0" plug0.alias = "Plug0"
plug0.device_id = "bb:bb:cc:dd:ee:ff_PLUG0DEVICEID" plug0.device_id = "bb:bb:cc:dd:ee:ff_PLUG0DEVICEID"
plug0.mac = "bb:bb:cc:dd:ee:ff" plug0.mac = "bb:bb:cc:dd:ee:ff"
plug0.protocol = _mock_protocol()
plug1 = _mocked_plug() plug1 = _mocked_plug()
plug1.device_id = "cc:bb:cc:dd:ee:ff_PLUG1DEVICEID" plug1.device_id = "cc:bb:cc:dd:ee:ff_PLUG1DEVICEID"
plug1.mac = "cc:bb:cc:dd:ee:ff" plug1.mac = "cc:bb:cc:dd:ee:ff"
plug1.alias = "Plug1" plug1.alias = "Plug1"
plug1.protocol = _mock_protocol()
strip.children = [plug0, plug1] strip.children = [plug0, plug1]
return strip return strip