diff --git a/devtools/dump_devinfo.py b/devtools/dump_devinfo.py index c1436aa1..a227a50d 100644 --- a/devtools/dump_devinfo.py +++ b/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: diff --git a/devtools/helpers/smartrequests.py b/devtools/helpers/smartrequests.py index e4941713..de0a53ff 100644 --- a/devtools/helpers/smartrequests.py +++ b/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, - "synchronize": None, # sync_env - "sunrise_sunset": None, # for schedules + "account": [], + "synchronize": [], # sync_env + "sunrise_sunset": [], # for schedules "led": [SmartRequest.get_led_info()], "cloud_connect": [SmartRequest.get_raw_request("get_connect_cloud_state")], - "iot_cloud": None, - "device_local_time": None, - "default_states": None, # in device_info + "iot_cloud": [], + "device_local_time": [], + "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 - "matter": None, + "current_protection": [], # overcurrent in device_info + "matter": [], "preset": [SmartRequest.get_preset_rules()], - "brightness": None, # in device_info - "color": None, # in device_info - "color_temperature": None, # in device_info + "brightness": [], # in device_info + "color": [], # 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")], } diff --git a/kasa/aestransport.py b/kasa/aestransport.py index bc1eacff..bbcc511f 100644 --- a/kasa/aestransport.py +++ b/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( diff --git a/kasa/cli.py b/kasa/cli.py index 9d4991aa..e922ec81 100755 --- a/kasa/cli.py +++ b/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) diff --git a/kasa/device_factory.py b/kasa/device_factory.py index 28a5e3b2..3550539c 100755 --- 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 @@ 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 diff --git a/kasa/smart/smartdevice.py b/kasa/smart/smartdevice.py index 04b84c38..dde8634f 100644 --- a/kasa/smart/smartdevice.py +++ b/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", {}) diff --git a/kasa/smartprotocol.py b/kasa/smartprotocol.py index 74f2275d..f61bac20 100644 --- a/kasa/smartprotocol.py +++ b/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() diff --git a/kasa/tests/fixtures/smart/P100_1.0.0_1.1.3.json b/kasa/tests/fixtures/smart/P100_1.0.0_1.1.3.json new file mode 100644 index 00000000..337c6f2c --- /dev/null +++ b/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" + } + } +} diff --git a/kasa/tests/test_aestransport.py b/kasa/tests/test_aestransport.py index a692ba9b..51f1e3d9 100644 --- a/kasa/tests/test_aestransport.py +++ b/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_dict = json_loads(response_data.decode()) - encrypted_response = self.encryption_session.encrypt( - json_dumps(decrypted_response_dict).encode() + decrypted_response_data = await decrypted_response.read() + encrypted_response = self.encryption_session.encrypt(decrypted_response_data) + response = ( + 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) diff --git a/kasa/tests/test_cli.py b/kasa/tests/test_cli.py index f44a924a..51155f40 100644 --- a/kasa/tests/test_cli.py +++ b/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"]) diff --git a/kasa/tests/test_device_factory.py b/kasa/tests/test_device_factory.py index 67ab39d5..7369a987 100644 --- a/kasa/tests/test_device_factory.py +++ b/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):