From bc59387437cea91d675584bb9057338aa5992817 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Tue, 28 Sep 2021 11:36:45 -0500
Subject: [PATCH] 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
---
 homeassistant/components/tplink/__init__.py |  2 ++
 tests/components/tplink/__init__.py         | 12 ++++++++++++
 2 files changed, 14 insertions(+)

diff --git a/homeassistant/components/tplink/__init__.py b/homeassistant/components/tplink/__init__.py
index e365f0f2453..5c40f61e2df 100644
--- a/homeassistant/components/tplink/__init__.py
+++ b/homeassistant/components/tplink/__init__.py
@@ -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
 
 
diff --git a/tests/components/tplink/__init__.py b/tests/components/tplink/__init__.py
index 1507685245b..870e05e970b 100644
--- a/tests/components/tplink/__init__.py
+++ b/tests/components/tplink/__init__.py
@@ -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