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]
if entry.entry_id not in hass_data:
return True
device: SmartDevice = hass.data[DOMAIN][entry.entry_id].device
if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
hass_data.pop(entry.entry_id)
await device.protocol.close()
return unload_ok

View File

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