Mirrored_Repos/python-kasa
1
0
Fork 0
mirror of https://github.com/python-kasa/python-kasa.git synced 2025-08-09 20:24:02 +00:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Cleanup discovery & add tests (#212)

Browse Source 
* Cleanup discovery & add tests

* discovered_devices_raw is not anymore available, as that can be accessed directly from the device objects
* test most of the discovery code paths
* some minor cleanups to test handling
* update discovery docs

* Move category check to be after the definitions

* skip a couple of tests requiring asyncmock not available on py37

* Remove return_raw usage from cli.discover
This commit is contained in:
Teemu R
2021-09-24 17:18:11 +02:00
committed by GitHub
parent bdb07a749c
commit acb221b1e0
4 changed files with 114 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files

56
kasa/discover.py
View File

@@ -3,7 +3,7 @@ import asyncio
import json
import logging
import socket
from typing import Awaitable, Callable, Dict, Mapping, Optional, Type, Union, cast
from typing import Awaitable, Callable, Dict, Optional, Type, cast
from kasa.protocol import TPLinkSmartHomeProtocol
from kasa.smartbulb import SmartBulb
@@ -17,6 +17,7 @@ _LOGGER = logging.getLogger(__name__)
OnDiscoveredCallable = Callable[[SmartDevice], Awaitable[None]]
DeviceDict = Dict[str, SmartDevice]
class _DiscoverProtocol(asyncio.DatagramProtocol):
@@ -25,8 +26,7 @@ class _DiscoverProtocol(asyncio.DatagramProtocol):
This is internal class, use :func:`Discover.discover`: instead.
"""
discovered_devices: Dict[str, SmartDevice]
discovered_devices_raw: Dict[str, Dict]
discovered_devices: DeviceDict
def __init__(
self,
@@ -43,7 +43,6 @@ class _DiscoverProtocol(asyncio.DatagramProtocol):
self.protocol = TPLinkSmartHomeProtocol()
self.target = (target, Discover.DISCOVERY_PORT)
self.discovered_devices = {}
self.discovered_devices_raw = {}
def connection_made(self, transport) -> None:
"""Set socket options for broadcasting."""
@@ -80,13 +79,9 @@ class _DiscoverProtocol(asyncio.DatagramProtocol):
device.update_from_discover_info(info)
self.discovered_devices[ip] = device
self.discovered_devices_raw[ip] = info
if device_class is not None:
if self.on_discovered is not None:
asyncio.ensure_future(self.on_discovered(device))
else:
_LOGGER.error("Received invalid response: %s", info)
if self.on_discovered is not None:
asyncio.ensure_future(self.on_discovered(device))
def error_received(self, ex):
"""Handle asyncio.Protocol errors."""
@@ -144,9 +139,8 @@ class Discover:
on_discovered=None,
timeout=5,
discovery_packets=3,
return_raw=False,
interface=None,
) -> Mapping[str, Union[SmartDevice, Dict]]:
) -> DeviceDict:
"""Discover supported devices.
Sends discovery message to 255.255.255.255:9999 in order
@@ -154,17 +148,17 @@ class Discover:
and waits for given timeout for answers from devices.
If you have multiple interfaces, you can use target parameter to specify the network for discovery.
If given, `on_discovered` coroutine will get passed with the :class:`SmartDevice`-derived object as parameter.
If given, `on_discovered` coroutine will get awaited with a :class:`SmartDevice`-derived object as parameter.
The results of the discovery are returned either as a list of :class:`SmartDevice`-derived objects
or as raw response dictionaries objects (if `return_raw` is True).
The results of the discovery are returned as a dict of :class:`SmartDevice`-derived objects keyed with IP addresses.
The devices are already initialized and all but emeter-related properties can be accessed directly.
:param target: The target address where to send the broadcast discovery queries if multi-homing (e.g. 192.168.xxx.255).
:param on_discovered: coroutine to execute on discovery
:param timeout: How long to wait for responses, defaults to 5
:param discovery_packets: Number of discovery packets are broadcasted.
:param return_raw: True to return JSON objects instead of Devices.
:return:
:param discovery_packets: Number of discovery packets to broadcast
:param interface: Bind to specific interface
:return: dictionary with discovered devices
"""
loop = asyncio.get_event_loop()
transport, protocol = await loop.create_datagram_endpoint(
@@ -186,9 +180,6 @@ class Discover:
_LOGGER.debug("Discovered %s devices", len(protocol.discovered_devices))
if return_raw:
return protocol.discovered_devices_raw
return protocol.discovered_devices
@staticmethod
@@ -204,12 +195,10 @@ class Discover:
info = await protocol.query(host, Discover.DISCOVERY_QUERY)
device_class = Discover._get_device_class(info)
if device_class is not None:
dev = device_class(host)
await dev.update()
return dev
dev = device_class(host)
await dev.update()
raise SmartDeviceException("Unable to discover device, received: %s" % info)
return dev
@staticmethod
def _get_device_class(info: dict) -> Type[SmartDevice]:
@@ -237,17 +226,4 @@ class Discover:
return SmartBulb
raise SmartDeviceException("Unknown device type: %s", type_)
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
loop = asyncio.get_event_loop()
async def _on_device(dev):
await dev.update()
_LOGGER.info("Got device: %s", dev)
devices = loop.run_until_complete(Discover.discover(on_discovered=_on_device))
for ip, dev in devices.items():
print(f"[{ip}] {dev}")
raise SmartDeviceException("Unknown device type: %s" % type_)