Merge remote-tracking branch 'upstream/master' into feat/dev_descriptors

devtools/dump_devinfo.py

@@ -12,6 +12,7 @@ import collections.abc
 import json
 import logging
 import re
+import traceback
 from collections import defaultdict, namedtuple
 from pathlib import Path
 from pprint import pprint
@@ -19,7 +20,7 @@ from typing import Dict, List, Union
 
 import asyncclick as click
 
-from devtools.helpers.smartrequests import COMPONENT_REQUESTS, SmartRequest
+from devtools.helpers.smartrequests import SmartRequest, get_component_requests
 from kasa import (
     AuthenticationException,
     Credentials,
@@ -35,6 +36,8 @@ from kasa.smart import SmartDevice
 Call = namedtuple("Call", "module method")
 SmartCall = namedtuple("SmartCall", "module request should_succeed")
 
+_LOGGER = logging.getLogger(__name__)
+
 
 def scrub(res):
     """Remove identifiers from the given dict."""
@@ -228,6 +231,8 @@ async def get_legacy_fixture(device):
             else:
                 click.echo(click.style("OK", fg="green"))
                 successes.append((test_call, info))
+        finally:
+            await device.protocol.close()
 
     final_query = defaultdict(defaultdict)
     final = defaultdict(defaultdict)
@@ -241,7 +246,8 @@ async def get_legacy_fixture(device):
         final = await device.protocol.query(final_query)
     except Exception as ex:
         _echo_error(f"Unable to query all successes at once: {ex}", bold=True, fg="red")
-
+    finally:
+        await device.protocol.close()
     if device._discovery_info and not device._discovery_info.get("system"):
         # Need to recreate a DiscoverResult here because we don't want the aliases
         # in the fixture, we want the actual field names as returned by the device.
@@ -316,7 +322,11 @@ async def _make_requests_or_exit(
         _echo_error(
             f"Unexpected exception querying {name} at once: {ex}",
         )
+        if _LOGGER.isEnabledFor(logging.DEBUG):
+            traceback.print_stack()
         exit(1)
+    finally:
+        await device.protocol.close()
 
 
 async def get_smart_fixture(device: SmartDevice, batch_size: int):
@@ -367,14 +377,15 @@ async def get_smart_fixture(device: SmartDevice, batch_size: int):
 
     for item in component_info_response["component_list"]:
         component_id = item["id"]
-        if requests := COMPONENT_REQUESTS.get(component_id):
+        ver_code = item["ver_code"]
+        if (requests := get_component_requests(component_id, ver_code)) is not None:
             component_test_calls = [
                 SmartCall(module=component_id, request=request, should_succeed=True)
                 for request in requests
             ]
             test_calls.extend(component_test_calls)
             should_succeed.extend(component_test_calls)
-        elif component_id not in COMPONENT_REQUESTS:
+        else:
             click.echo(f"Skipping {component_id}..", nl=False)
             click.echo(click.style("UNSUPPORTED", fg="yellow"))
 
@@ -396,7 +407,11 @@ async def get_smart_fixture(device: SmartDevice, batch_size: int):
             if (
                 not test_call.should_succeed
                 and hasattr(ex, "error_code")
-                and ex.error_code == SmartErrorCode.UNKNOWN_METHOD_ERROR
+                and ex.error_code
+                in [
+                    SmartErrorCode.UNKNOWN_METHOD_ERROR,
+                    SmartErrorCode.TRANSPORT_NOT_AVAILABLE_ERROR,
+                ]
             ):
                 click.echo(click.style("FAIL - EXPECTED", fg="green"))
             else:
@@ -410,6 +425,8 @@ async def get_smart_fixture(device: SmartDevice, batch_size: int):
                 else:
                     click.echo(click.style("OK", fg="green"))
                 successes.append(test_call)
+        finally:
+            await device.protocol.close()
 
     requests = []
     for succ in successes:

devtools/helpers/smartrequests.py

@@ -133,11 +133,14 @@ class SmartRequest:
         return SmartRequest("get_device_usage")
 
     @staticmethod
-    def device_info_list() -> List["SmartRequest"]:
+    def device_info_list(ver_code) -> List["SmartRequest"]:
         """Get device info list."""
+        if ver_code == 1:
+            return [SmartRequest.get_device_info()]
         return [
             SmartRequest.get_device_info(),
             SmartRequest.get_device_usage(),
+            SmartRequest.get_auto_update_info(),
         ]
 
     @staticmethod
@@ -149,7 +152,6 @@ class SmartRequest:
     def firmware_info_list() -> List["SmartRequest"]:
         """Get info list."""
         return [
-            SmartRequest.get_auto_update_info(),
             SmartRequest.get_raw_request("get_fw_download_state"),
             SmartRequest.get_raw_request("get_latest_fw"),
         ]
@@ -165,9 +167,13 @@ class SmartRequest:
         return SmartRequest("get_device_time")
 
     @staticmethod
-    def get_wireless_scan_info() -> "SmartRequest":
+    def get_wireless_scan_info(
+        params: Optional[GetRulesParams] = None
+    ) -> "SmartRequest":
         """Get wireless scan info."""
-        return SmartRequest("get_wireless_scan_info")
+        return SmartRequest(
+            "get_wireless_scan_info", params or SmartRequest.GetRulesParams()
+        )
 
     @staticmethod
     def get_schedule_rules(params: Optional[GetRulesParams] = None) -> "SmartRequest":
@@ -294,9 +300,13 @@ class SmartRequest:
     @staticmethod
     def get_component_info_requests(component_nego_response) -> List["SmartRequest"]:
         """Get a list of requests based on the component info response."""
-        request_list = []
+        request_list: List["SmartRequest"] = []
         for component in component_nego_response["component_list"]:
-            if requests := COMPONENT_REQUESTS.get(component["id"]):
+            if (
+                requests := get_component_requests(
+                    component["id"], int(component["ver_code"])
+                )
+            ) is not None:
                 request_list.extend(requests)
         return request_list
 
@@ -314,8 +324,17 @@ class SmartRequest:
         return request
 
 
+def get_component_requests(component_id, ver_code):
+    """Get the requests supported by the component and version."""
+    if (cr := COMPONENT_REQUESTS.get(component_id)) is None:
+        return None
+    if callable(cr):
+        return cr(ver_code)
+    return cr
+
+
 COMPONENT_REQUESTS = {
-    "device": SmartRequest.device_info_list(),
+    "device": SmartRequest.device_info_list,
     "firmware": SmartRequest.firmware_info_list(),
     "quick_setup": [SmartRequest.qs_component_nego()],
     "inherit": [SmartRequest.get_raw_request("get_inherit_info")],
@@ -324,33 +343,33 @@ COMPONENT_REQUESTS = {
     "schedule": SmartRequest.schedule_info_list(),
     "countdown": [SmartRequest.get_countdown_rules()],
     "antitheft": [SmartRequest.get_antitheft_rules()],
-    "account": None,
+    "account": [],
-    "synchronize": None,  # sync_env
+    "synchronize": [],  # sync_env
-    "sunrise_sunset": None,  # for schedules
+    "sunrise_sunset": [],  # for schedules
     "led": [SmartRequest.get_led_info()],
     "cloud_connect": [SmartRequest.get_raw_request("get_connect_cloud_state")],
-    "iot_cloud": None,
+    "iot_cloud": [],
-    "device_local_time": None,
+    "device_local_time": [],
-    "default_states": None,  # in device_info
+    "default_states": [],  # in device_info
     "auto_off": [SmartRequest.get_auto_off_config()],
-    "localSmart": None,
+    "localSmart": [],
     "energy_monitoring": SmartRequest.energy_monitoring_list(),
     "power_protection": SmartRequest.power_protection_list(),
-    "current_protection": None,  # overcurrent in device_info
+    "current_protection": [],  # overcurrent in device_info
-    "matter": None,
+    "matter": [],
     "preset": [SmartRequest.get_preset_rules()],
-    "brightness": None,  # in device_info
+    "brightness": [],  # in device_info
-    "color": None,  # in device_info
+    "color": [],  # in device_info
-    "color_temperature": None,  # in device_info
+    "color_temperature": [],  # in device_info
     "auto_light": [SmartRequest.get_auto_light_info()],
     "light_effect": [SmartRequest.get_dynamic_light_effect_rules()],
-    "bulb_quick_control": None,
+    "bulb_quick_control": [],
     "on_off_gradually": [SmartRequest.get_raw_request("get_on_off_gradually_info")],
-    "light_strip": None,
+    "light_strip": [],
     "light_strip_lighting_effect": [
         SmartRequest.get_raw_request("get_lighting_effect")
     ],
-    "music_rhythm": None,  # music_rhythm_enable in device_info
+    "music_rhythm": [],  # music_rhythm_enable in device_info
     "segment": [SmartRequest.get_raw_request("get_device_segment")],
     "segment_effect": [SmartRequest.get_raw_request("get_segment_effect_rule")],
 }

kasa/aestransport.py

@@ -3,7 +3,7 @@
 Based on the work of https://github.com/petretiandrea/plugp100
 under compatible GNU GPL3 license.
 """
-
+import asyncio
 import base64
 import hashlib
 import logging
@@ -39,6 +39,7 @@ _LOGGER = logging.getLogger(__name__)
 
 ONE_DAY_SECONDS = 86400
 SESSION_EXPIRE_BUFFER_SECONDS = 60 * 20
+BACKOFF_SECONDS_AFTER_LOGIN_ERROR = 1
 
 
 def _sha1(payload: bytes) -> str:
@@ -184,8 +185,24 @@ class AesTransport(BaseTransport):
             assert self._encryption_session is not None
 
         raw_response: str = resp_dict["result"]["response"]
-        response = self._encryption_session.decrypt(raw_response.encode())
+
-        return json_loads(response)  # type: ignore[return-value]
+        try:
+            response = self._encryption_session.decrypt(raw_response.encode())
+            ret_val = json_loads(response)
+        except Exception as ex:
+            try:
+                ret_val = json_loads(raw_response)
+                _LOGGER.debug(
+                    "Received unencrypted response over secure passthrough from %s",
+                    self._host,
+                )
+            except Exception:
+                raise SmartDeviceException(
+                    f"Unable to decrypt response from {self._host}, "
+                    + f"error: {ex}, response: {raw_response}",
+                    ex,
+                ) from ex
+        return ret_val  # type: ignore[return-value]
 
     async def perform_login(self):
         """Login to the device."""
@@ -199,6 +216,7 @@ class AesTransport(BaseTransport):
                     self._default_credentials = get_default_credentials(
                         DEFAULT_CREDENTIALS["TAPO"]
                     )
+                    await asyncio.sleep(BACKOFF_SECONDS_AFTER_LOGIN_ERROR)
                 await self.perform_handshake()
                 await self.try_login(self._get_login_params(self._default_credentials))
                 _LOGGER.debug(

kasa/cli.py

@@ -5,6 +5,7 @@ import json
 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
@@ -217,7 +218,7 @@ def json_formatter_cb(result, **kwargs):
 @click.option(
     "--discovery-timeout",
     envvar="KASA_DISCOVERY_TIMEOUT",
-    default=3,
+    default=5,
     required=False,
     show_default=True,
     help="Timeout for discovery.",
@@ -349,11 +350,16 @@ async def cli(
         )
         dev = await Device.connect(config=config)
     else:
-        echo("No --type or --device-family and --encrypt-type defined, discovering..")
+        echo(
+            "No --type or --device-family and --encrypt-type defined, "
+            + f"discovering for {discovery_timeout} seconds.."
+        )
         dev = await Discover.discover_single(
             host,
             port=port,
             credentials=credentials,
+            timeout=timeout,
+            discovery_timeout=discovery_timeout,
         )
 
     # Skip update on specific commands, or if device factory,
@@ -361,7 +367,14 @@ async def cli(
     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)

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 @@ async def connect(*, host: Optional[str] = None, config: DeviceConfig) -> "Devic
             )
             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
 

kasa/smart/smartdevice.py

@@ -70,7 +70,7 @@ class SmartDevice(Device):
             resp = await self.protocol.query("component_nego")
             self._components_raw = resp["component_nego"]
             self._components = {
-                comp["id"]: comp["ver_code"]
+                comp["id"]: int(comp["ver_code"])
                 for comp in self._components_raw["component_list"]
             }
             await self._initialize_modules()
@@ -87,9 +87,14 @@ class SmartDevice(Device):
                 "get_current_power": None,
             }
 
+        if self._components["device"] >= 2:
+            extra_reqs = {
+                **extra_reqs,
+                "get_device_usage": None,
+            }
+
         req = {
             "get_device_info": None,
-            "get_device_usage": None,
             "get_device_time": None,
             **extra_reqs,
         }
@@ -97,8 +102,9 @@ class SmartDevice(Device):
         resp = await self.protocol.query(req)
 
         self._info = resp["get_device_info"]
-        self._usage = resp["get_device_usage"]
         self._time = resp["get_device_time"]
+        # Device usage is not available on older firmware versions
+        self._usage = resp.get("get_device_usage", {})
         # Emeter is not always available, but we set them still for now.
         self._energy = resp.get("get_energy_usage", {})
         self._emeter = resp.get("get_current_power", {})

kasa/smartprotocol.py

@@ -82,6 +82,7 @@ class SmartProtocol(BaseProtocol):
                 if retry >= retry_count:
                     _LOGGER.debug("Giving up on %s after %s retries", self._host, retry)
                     raise ex
+                await asyncio.sleep(self.BACKOFF_SECONDS_AFTER_TIMEOUT)
                 continue
             except TimeoutException as ex:
                 await self._transport.reset()

kasa/tests/fixtures/smart/P100_1.0.0_1.1.3.json

@@ -0,0 +1,173 @@
+{
+    "component_nego": {
+        "component_list": [
+            {
+                "id": "device",
+                "ver_code": 1
+            },
+            {
+                "id": "firmware",
+                "ver_code": 1
+            },
+            {
+                "id": "quick_setup",
+                "ver_code": 1
+            },
+            {
+                "id": "time",
+                "ver_code": 1
+            },
+            {
+                "id": "wireless",
+                "ver_code": 1
+            },
+            {
+                "id": "schedule",
+                "ver_code": 1
+            },
+            {
+                "id": "countdown",
+                "ver_code": 1
+            },
+            {
+                "id": "antitheft",
+                "ver_code": 1
+            },
+            {
+                "id": "account",
+                "ver_code": 1
+            },
+            {
+                "id": "synchronize",
+                "ver_code": 1
+            },
+            {
+                "id": "sunrise_sunset",
+                "ver_code": 1
+            },
+            {
+                "id": "led",
+                "ver_code": 1
+            },
+            {
+                "id": "cloud_connect",
+                "ver_code": 1
+            }
+        ]
+    },
+    "discovery_result": {
+        "device_id": "00000000000000000000000000000000",
+        "device_model": "P100",
+        "device_type": "SMART.TAPOPLUG",
+        "factory_default": false,
+        "ip": "127.0.0.123",
+        "mac": "1C-3B-F3-00-00-00",
+        "mgt_encrypt_schm": {
+            "encrypt_type": "AES",
+            "http_port": 80,
+            "is_support_https": false
+        },
+        "owner": "00000000000000000000000000000000"
+    },
+    "get_antitheft_rules": {
+        "antitheft_rule_max_count": 1,
+        "enable": false,
+        "rule_list": []
+    },
+    "get_connect_cloud_state": {
+        "status": -1001
+    },
+    "get_countdown_rules": {
+        "countdown_rule_max_count": 1,
+        "enable": false,
+        "rule_list": []
+    },
+    "get_device_info": {
+        "avatar": "plug",
+        "device_id": "0000000000000000000000000000000000000000",
+        "device_on": true,
+        "fw_id": "00000000000000000000000000000000",
+        "fw_ver": "1.1.3 Build 20191017 Rel. 57937",
+        "has_set_location_info": true,
+        "hw_id": "00000000000000000000000000000000",
+        "hw_ver": "1.0.0",
+        "ip": "127.0.0.123",
+        "latitude": 0,
+        "location": "hallway",
+        "longitude": 0,
+        "mac": "1C-3B-F3-00-00-00",
+        "model": "P100",
+        "nickname": "I01BU0tFRF9OQU1FIw==",
+        "oem_id": "00000000000000000000000000000000",
+        "on_time": 6868,
+        "overheated": false,
+        "signal_level": 2,
+        "specs": "US",
+        "ssid": "I01BU0tFRF9TU0lEIw==",
+        "time_usage_past30": 114,
+        "time_usage_past7": 114,
+        "time_usage_today": 114,
+        "type": "SMART.TAPOPLUG"
+    },
+    "get_device_time": {
+        "region": "Europe/London",
+        "time_diff": 0,
+        "timestamp": 1707905077
+    },
+    "get_fw_download_state": {
+        "download_progress": 0,
+        "reboot_time": 10,
+        "status": 0,
+        "upgrade_time": 0
+    },
+    "get_latest_fw": {
+        "fw_size": 786432,
+        "fw_ver": "1.3.7 Build 20230711 Rel.61904",
+        "hw_id": "00000000000000000000000000000000",
+        "need_to_upgrade": true,
+        "oem_id": "00000000000000000000000000000000",
+        "release_date": "2023-07-26",
+        "release_note": "Modifications and Bug fixes:\nEnhanced device security.",
+        "type": 3
+    },
+    "get_led_info": {
+        "led_rule": "always",
+        "led_status": true
+    },
+    "get_next_event": {
+        "action": -1,
+        "e_time": 0,
+        "id": "0",
+        "s_time": 0,
+        "type": 0
+    },
+    "get_schedule_rules": {
+        "enable": false,
+        "rule_list": [],
+        "schedule_rule_max_count": 20,
+        "start_index": 0,
+        "sum": 0
+    },
+    "get_wireless_scan_info": {
+        "ap_list": [],
+        "start_index": 0,
+        "sum": 0,
+        "wep_supported": false
+    },
+    "qs_component_nego": {
+        "component_list": [
+            {
+                "id": "quick_setup",
+                "ver_code": 1
+            },
+            {
+                "id": "sunrise_sunset",
+                "ver_code": 1
+            }
+        ],
+        "extra_info": {
+            "device_model": "P100",
+            "device_type": "SMART.TAPOPLUG"
+        }
+    }
+}

kasa/tests/test_aestransport.py

@@ -1,5 +1,6 @@
 import base64
 import json
+import logging
 import random
 import string
 import time
@@ -180,6 +181,67 @@ async def test_send(mocker, status_code, error_code, inner_error_code, expectati
         assert "result" in res
 
 
+async def test_unencrypted_response(mocker, caplog):
+    host = "127.0.0.1"
+    mock_aes_device = MockAesDevice(host, 200, 0, 0, do_not_encrypt_response=True)
+    mocker.patch.object(aiohttp.ClientSession, "post", side_effect=mock_aes_device.post)
+
+    transport = AesTransport(
+        config=DeviceConfig(host, credentials=Credentials("foo", "bar"))
+    )
+    transport._state = TransportState.ESTABLISHED
+    transport._session_expire_at = time.time() + 86400
+    transport._encryption_session = mock_aes_device.encryption_session
+    transport._token_url = transport._app_url.with_query(
+        f"token={mock_aes_device.token}"
+    )
+
+    request = {
+        "method": "get_device_info",
+        "params": None,
+        "request_time_milis": round(time.time() * 1000),
+        "requestID": 1,
+        "terminal_uuid": "foobar",
+    }
+    caplog.set_level(logging.DEBUG)
+    res = await transport.send(json_dumps(request))
+    assert "result" in res
+    assert (
+        "Received unencrypted response over secure passthrough from 127.0.0.1"
+        in caplog.text
+    )
+
+
+async def test_unencrypted_response_invalid_json(mocker, caplog):
+    host = "127.0.0.1"
+    mock_aes_device = MockAesDevice(
+        host, 200, 0, 0, do_not_encrypt_response=True, send_response=b"Foobar"
+    )
+    mocker.patch.object(aiohttp.ClientSession, "post", side_effect=mock_aes_device.post)
+
+    transport = AesTransport(
+        config=DeviceConfig(host, credentials=Credentials("foo", "bar"))
+    )
+    transport._state = TransportState.ESTABLISHED
+    transport._session_expire_at = time.time() + 86400
+    transport._encryption_session = mock_aes_device.encryption_session
+    transport._token_url = transport._app_url.with_query(
+        f"token={mock_aes_device.token}"
+    )
+
+    request = {
+        "method": "get_device_info",
+        "params": None,
+        "request_time_milis": round(time.time() * 1000),
+        "requestID": 1,
+        "terminal_uuid": "foobar",
+    }
+    caplog.set_level(logging.DEBUG)
+    msg = f"Unable to decrypt response from {host}, error: Incorrect padding, response: Foobar"
+    with pytest.raises(SmartDeviceException, match=msg):
+        await transport.send(json_dumps(request))
+
+
 ERRORS = [e for e in SmartErrorCode if e != 0]
 
 
@@ -233,15 +295,28 @@ class MockAesDevice:
             pass
 
         async def read(self):
-            return json_dumps(self._json).encode()
+            if isinstance(self._json, dict):
+                return json_dumps(self._json).encode()
+            return self._json
 
     encryption_session = AesEncyptionSession(KEY_IV[:16], KEY_IV[16:])
 
-    def __init__(self, host, status_code=200, error_code=0, inner_error_code=0):
+    def __init__(
+        self,
+        host,
+        status_code=200,
+        error_code=0,
+        inner_error_code=0,
+        *,
+        do_not_encrypt_response=False,
+        send_response=None,
+    ):
         self.host = host
         self.status_code = status_code
         self.error_code = error_code
         self._inner_error_code = inner_error_code
+        self.do_not_encrypt_response = do_not_encrypt_response
+        self.send_response = send_response
         self.http_client = HttpClient(DeviceConfig(self.host))
         self.inner_call_count = 0
         self.token = "".join(random.choices(string.ascii_uppercase, k=32))  # noqa: S311
@@ -289,13 +364,15 @@ class MockAesDevice:
         decrypted_request_dict = json_loads(decrypted_request)
         decrypted_response = await self._post(url, decrypted_request_dict)
         async with decrypted_response:
-            response_data = await decrypted_response.read()
+            decrypted_response_data = await decrypted_response.read()
-            decrypted_response_dict = json_loads(response_data.decode())
+        encrypted_response = self.encryption_session.encrypt(decrypted_response_data)
-        encrypted_response = self.encryption_session.encrypt(
+        response = (
-            json_dumps(decrypted_response_dict).encode()
+            decrypted_response_data
+            if self.do_not_encrypt_response
+            else encrypted_response
         )
         result = {
-            "result": {"response": encrypted_response.decode()},
+            "result": {"response": response.decode()},
             "error_code": self.error_code,
         }
         return self._mock_response(self.status_code, result)
@@ -310,5 +387,6 @@ class MockAesDevice:
 
     async def _return_send_response(self, url: URL, json: Dict[str, Any]):
         result = {"result": {"method": None}, "error_code": self.inner_error_code}
+        response = self.send_response if self.send_response else result
         self.inner_call_count += 1
-        return self._mock_response(self.status_code, result)
+        return self._mock_response(self.status_code, response)

kasa/tests/test_cli.py

@@ -7,6 +7,7 @@ from asyncclick.testing import CliRunner
 
 from kasa import (
     AuthenticationException,
+    Credentials,
     Device,
     EmeterStatus,
     SmartDeviceException,
@@ -351,7 +352,9 @@ async def test_credentials(discovery_mock, mocker):
 async def test_without_device_type(dev, mocker):
     """Test connecting without the device type."""
     runner = CliRunner()
-    mocker.patch("kasa.discover.Discover.discover_single", return_value=dev)
+    discovery_mock = mocker.patch(
+        "kasa.discover.Discover.discover_single", return_value=dev
+    )
     # These will mock the features to avoid accessing non-existing
     mocker.patch("kasa.device.Device.features", return_value={})
     mocker.patch("kasa.iot.iotdevice.IotDevice.features", return_value={})
@@ -365,9 +368,18 @@ async def test_without_device_type(dev, mocker):
             "foo",
             "--password",
             "bar",
+            "--discovery-timeout",
+            "7",
         ],
     )
     assert res.exit_code == 0
+    discovery_mock.assert_called_once_with(
+        "127.0.0.1",
+        port=None,
+        credentials=Credentials("foo", "bar"),
+        timeout=5,
+        discovery_timeout=7,
+    )
 
 
 @pytest.mark.parametrize("auth_param", ["--username", "--password"])

kasa/tests/test_device_factory.py

@@ -53,7 +53,7 @@ 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,
     )
@@ -61,8 +61,9 @@ async def test_connect(
     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):