Ensure connections are closed when cli is finished (#752)

sdb9696 · web-flow · commit 45f251e57e03 · 2024-02-14T17:03:50.000Z
* Ensure connections are closed when cli is finished

* Test for close calls on error and success
diff --git a/kasa/cli.py b/kasa/cli.py
@@ -5,6 +5,7 @@
 import logging
 import re
 import sys
+from contextlib import asynccontextmanager
 from functools import singledispatch, wraps
 from pprint import pformat as pf
 from typing import Any, Dict, cast
@@ -365,7 +366,14 @@ def _nop_echo(*args, **kwargs):
     if ctx.invoked_subcommand not in SKIP_UPDATE_COMMANDS and not device_family:
         await dev.update()
 
-    ctx.obj = dev
+    @asynccontextmanager
+    async def async_wrapped_device(device: Device):
+        try:
+            yield device
+        finally:
+            await device.disconnect()
+
+    ctx.obj = await ctx.with_async_resource(async_wrapped_device(dev))
 
     if ctx.invoked_subcommand is None:
         return await ctx.invoke(state)
diff --git a/kasa/device_factory.py b/kasa/device_factory.py
@@ -49,6 +49,20 @@ async def connect(*, host: Optional[str] = None, config: DeviceConfig) -> "Devic
     if host:
         config = DeviceConfig(host=host)
 
+    if (protocol := get_protocol(config=config)) is None:
+        raise UnsupportedDeviceException(
+            f"Unsupported device for {config.host}: "
+            + f"{config.connection_type.device_family.value}"
+        )
+
+    try:
+        return await _connect(config, protocol)
+    except:
+        await protocol.close()
+        raise
+
+
+async def _connect(config: DeviceConfig, protocol: BaseProtocol) -> "Device":
     debug_enabled = _LOGGER.isEnabledFor(logging.DEBUG)
     if debug_enabled:
         start_time = time.perf_counter()
@@ -63,12 +77,6 @@ def _perf_log(has_params, perf_type):
             )
             start_time = time.perf_counter()
 
-    if (protocol := get_protocol(config=config)) is None:
-        raise UnsupportedDeviceException(
-            f"Unsupported device for {config.host}: "
-            + f"{config.connection_type.device_family.value}"
-        )
-
     device_class: Optional[Type[Device]]
     device: Optional[Device] = None
 
diff --git a/kasa/tests/test_device_factory.py b/kasa/tests/test_device_factory.py
@@ -53,16 +53,17 @@ async def test_connect(
         host=host, credentials=Credentials("foor", "bar"), connection_type=ctype
     )
     protocol_class = get_protocol(config).__class__
-
+    close_mock = mocker.patch.object(protocol_class, "close")
     dev = await connect(
         config=config,
     )
     assert isinstance(dev, device_class)
     assert isinstance(dev.protocol, protocol_class)
 
     assert dev.config == config
-
+    assert close_mock.call_count == 0
     await dev.disconnect()
+    assert close_mock.call_count == 1
 
 
 @pytest.mark.parametrize("custom_port", [123, None])
@@ -116,8 +117,12 @@ async def test_connect_query_fails(all_fixture_data: dict, mocker):
     config = DeviceConfig(
         host=host, credentials=Credentials("foor", "bar"), connection_type=ctype
     )
+    protocol_class = get_protocol(config).__class__
+    close_mock = mocker.patch.object(protocol_class, "close")
+    assert close_mock.call_count == 0
     with pytest.raises(SmartDeviceException):
         await connect(config=config)
+    assert close_mock.call_count == 1
 
 
 async def test_connect_http_client(all_fixture_data, mocker):
