Skip to content

Re-add protocol_class parameter to connect #551

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.

Sign up for GitHub

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
Nov 28, 2023
Merged

Re-add protocol_class parameter to connect #551

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
201c1bb
Re-add protocol_class parameter to connect
sdb9696 Nov 22, 2023
d0e4cc5
Merge branch 'master' into add_protocol_to_connect
bdraco Nov 28, 2023
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
8 changes: 8 additions & 0 deletions kasa/device_factory.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
from .credentials import Credentials
from .device_type import DeviceType
from .exceptions import UnsupportedDeviceException
from .protocol import TPLinkProtocol
from .smartbulb import SmartBulb
from .smartdevice import SmartDevice, SmartDeviceException
from .smartdimmer import SmartDimmer
Expand All @@ -32,6 +33,7 @@ async def connect(
timeout=5,
credentials: Optional[Credentials] = None,
device_type: Optional[DeviceType] = None,
protocol_class: Optional[Type[TPLinkProtocol]] = None,
) -> "SmartDevice":
"""Connect to a single device by the given IP address.

Expand All @@ -50,6 +52,8 @@ async def connect(
If not given, the device type is discovered by querying the device.
If the device type is already known, it is preferred to pass it
to avoid the extra query to the device to discover its type.
:param protocol_class: Optionally provide the protocol class
to use.
:rtype: SmartDevice
:return: Object for querying/controlling found device.
"""
Expand All @@ -62,6 +66,8 @@ async def connect(
dev: SmartDevice = klass(
host=host, port=port, credentials=credentials, timeout=timeout
)
if protocol_class is not None:
dev.protocol = protocol_class(host, credentials=credentials)
await dev.update()
if debug_enabled:
end_time = time.perf_counter()
Expand All @@ -76,6 +82,8 @@ async def connect(
unknown_dev = SmartDevice(
host=host, port=port, credentials=credentials, timeout=timeout
)
if protocol_class is not None:
unknown_dev.protocol = protocol_class(host, credentials=credentials)
await unknown_dev.update()
device_class = get_device_class_from_info(unknown_dev.internal_state)
dev = device_class(host=host, port=port, credentials=credentials, timeout=timeout)
Expand Down
27 changes: 27 additions & 0 deletions kasa/tests/test_device_factory.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
SmartPlug,
)
from kasa.device_factory import connect
from kasa.klapprotocol import TPLinkKlap
from kasa.protocol import TPLinkProtocol, TPLinkSmartHomeProtocol


@pytest.mark.parametrize("custom_port", [123, None])
Expand Down Expand Up @@ -72,3 +74,28 @@ async def test_connect_logs_connect_time(
logging.getLogger("kasa").setLevel(logging.DEBUG)
await connect(host)
assert "seconds to connect" in caplog.text


@pytest.mark.parametrize("device_type", [DeviceType.Plug, None])
@pytest.mark.parametrize(
("protocol_in", "protocol_result"),
(
(None, TPLinkSmartHomeProtocol),
(TPLinkKlap, TPLinkKlap),
(TPLinkSmartHomeProtocol, TPLinkSmartHomeProtocol),
),
)
async def test_connect_pass_protocol(
discovery_data: dict,
mocker,
device_type: DeviceType,
protocol_in: Type[TPLinkProtocol],
protocol_result: Type[TPLinkProtocol],
):
"""Test that if the protocol is passed in it's gets set correctly."""
host = "127.0.0.1"
mocker.patch("kasa.TPLinkSmartHomeProtocol.query", return_value=discovery_data)
mocker.patch("kasa.TPLinkKlap.query", return_value=discovery_data)

dev = await connect(host, device_type=device_type, protocol_class=protocol_in)
assert isinstance(dev.protocol, protocol_result)