8000 Ensure connections are closed when cli is finished by sdb9696 · Pull Request #752 · python-kasa/python-kasa · GitHub
[go: up one dir, main page]

Skip to content

Ensure connections are closed when cli is finished #752

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Feb 14, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 9 additions & 1 deletion kasa/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.inv 8000 oked_subcommand is None:
return await ctx.invoke(state)
Expand Down
20 changes: 14 additions & 6 deletions kasa/device_factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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

Expand Down
9 changes: 7 additions & 2 deletions kasa/tests/test_device_factory.py
8000
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand Down Expand Up @@ -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):
Expand Down
0