Add klap support for TAPO protocol by splitting out Transports and Protocols (#557)

* Add support for TAPO/SMART KLAP and seperate transports from protocols

* Add tests and some review changes

* Update following review

* Updates following review

kasa/tests/conftest.py

@@ -2,27 +2,45 @@ import asyncio
 import glob
 import json
 import os
+from dataclasses import dataclass
+from json import dumps as json_dumps
 from os.path import basename
 from pathlib import Path, PurePath
-from typing import Dict
+from typing import Dict, Optional
 from unittest.mock import MagicMock
 
 import pytest  # type: ignore # see https://github.com/pytest-dev/pytest/issues/3342
 
 from kasa import (
+    Credentials,
     Discover,
     SmartBulb,
     SmartDimmer,
     SmartLightStrip,
     SmartPlug,
     SmartStrip,
+    TPLinkSmartHomeProtocol,
 )
+from kasa.tapo import TapoDevice, TapoPlug
 
-from .newfakes import FakeTransportProtocol
+from .newfakes import FakeSmartProtocol, FakeTransportProtocol
 
-SUPPORTED_DEVICES = glob.glob(
-    os.path.dirname(os.path.abspath(__file__)) + "/fixtures/*.json"
-)
+SUPPORTED_IOT_DEVICES = [
+    (device, "IOT")
+    for device in glob.glob(
+        os.path.dirname(os.path.abspath(__file__)) + "/fixtures/*.json"
+    )
+]
+
+SUPPORTED_SMART_DEVICES = [
+    (device, "SMART")
+    for device in glob.glob(
+        os.path.dirname(os.path.abspath(__file__)) + "/fixtures/smart/*.json"
+    )
+]
+
+
+SUPPORTED_DEVICES = SUPPORTED_IOT_DEVICES + SUPPORTED_SMART_DEVICES
 
 
 LIGHT_STRIPS = {"KL400", "KL430", "KL420"}
@@ -55,43 +73,59 @@ PLUGS = {
     "KP401",
     "KS200M",
 }
+
 STRIPS = {"HS107", "HS300", "KP303", "KP200", "KP400", "EP40"}
 DIMMERS = {"ES20M", "HS220", "KS220M", "KS230", "KP405"}
 
 DIMMABLE = {*BULBS, *DIMMERS}
 WITH_EMETER = {"HS110", "HS300", "KP115", "KP125", *BULBS}
 
-ALL_DEVICES = BULBS.union(PLUGS).union(STRIPS).union(DIMMERS)
+ALL_DEVICES_IOT = BULBS.union(PLUGS).union(STRIPS).union(DIMMERS)
+
+PLUGS_SMART = {"P110"}
+ALL_DEVICES_SMART = PLUGS_SMART
+
+ALL_DEVICES = ALL_DEVICES_IOT.union(ALL_DEVICES_SMART)
 
 IP_MODEL_CACHE: Dict[str, str] = {}
 
 
-def filter_model(desc, filter):
-    filtered = list()
-    for dev in SUPPORTED_DEVICES:
-        for filt in filter:
-            if filt in basename(dev):
-                filtered.append(dev)
+def idgenerator(paramtuple):
+    return basename(paramtuple[0]) + (
+        "" if paramtuple[1] == "IOT" else "-" + paramtuple[1]
+    )
 
-    filtered_basenames = [basename(f) for f in filtered]
+
+def filter_model(desc, model_filter, protocol_filter=None):
+    if not protocol_filter:
+        protocol_filter = {"IOT"}
+    filtered = list()
+    for file, protocol in SUPPORTED_DEVICES:
+        if protocol in protocol_filter:
+            file_model = basename(file).split("_")[0]
+            for model in model_filter:
+                if model in file_model:
+                    filtered.append((file, protocol))
+
+    filtered_basenames = [basename(f) + "-" + p for f, p in filtered]
     print(f"{desc}: {filtered_basenames}")
     return filtered
 
 
-def parametrize(desc, devices, ids=None):
+def parametrize(desc, devices, protocol_filter=None, ids=None):
     return pytest.mark.parametrize(
-        "dev", filter_model(desc, devices), indirect=True, ids=ids
+        "dev", filter_model(desc, devices, protocol_filter), indirect=True, ids=ids
     )
 
 
 has_emeter = parametrize("has emeter", WITH_EMETER)
-no_emeter = parametrize("no emeter", ALL_DEVICES - WITH_EMETER)
+no_emeter = parametrize("no emeter", ALL_DEVICES_IOT - WITH_EMETER)
 
-bulb = parametrize("bulbs", BULBS, ids=basename)
-plug = parametrize("plugs", PLUGS, ids=basename)
-strip = parametrize("strips", STRIPS, ids=basename)
-dimmer = parametrize("dimmers", DIMMERS, ids=basename)
-lightstrip = parametrize("lightstrips", LIGHT_STRIPS, ids=basename)
+bulb = parametrize("bulbs", BULBS, ids=idgenerator)
+plug = parametrize("plugs", PLUGS, ids=idgenerator)
+strip = parametrize("strips", STRIPS, ids=idgenerator)
+dimmer = parametrize("dimmers", DIMMERS, ids=idgenerator)
+lightstrip = parametrize("lightstrips", LIGHT_STRIPS, ids=idgenerator)
 
 # bulb types
 dimmable = parametrize("dimmable", DIMMABLE)
@@ -101,6 +135,58 @@ non_variable_temp = parametrize("non-variable color temp", BULBS - VARIABLE_TEMP
 color_bulb = parametrize("color bulbs", COLOR_BULBS)
 non_color_bulb = parametrize("non-color bulbs", BULBS - COLOR_BULBS)
 
+plug_smart = parametrize(
+    "plug devices smart", PLUGS_SMART, protocol_filter={"SMART"}, ids=idgenerator
+)
+device_smart = parametrize(
+    "devices smart", ALL_DEVICES_SMART, protocol_filter={"SMART"}, ids=idgenerator
+)
+device_iot = parametrize(
+    "devices iot", ALL_DEVICES_IOT, protocol_filter={"IOT"}, ids=idgenerator
+)
+
+
+def get_fixture_data():
+    """Return raw discovery file contents as JSON. Used for discovery tests."""
+    fixture_data = {}
+    for file, protocol in SUPPORTED_DEVICES:
+        p = Path(file)
+        if not p.is_absolute():
+            folder = Path(__file__).parent / "fixtures"
+            if protocol == "SMART":
+                folder = folder / "smart"
+            p = folder / file
+
+        with open(p) as f:
+            fixture_data[basename(p)] = json.load(f)
+    return fixture_data
+
+
+FIXTURE_DATA = get_fixture_data()
+
+
+def filter_fixtures(desc, root_filter):
+    filtered = {}
+    for key, val in FIXTURE_DATA.items():
+        if root_filter in val:
+            filtered[key] = val
+
+    print(f"{desc}: {filtered.keys()}")
+    return filtered
+
+
+def parametrize_discovery(desc, root_key):
+    filtered_fixtures = filter_fixtures(desc, root_key)
+    return pytest.mark.parametrize(
+        "discovery_data",
+        filtered_fixtures.values(),
+        indirect=True,
+        ids=filtered_fixtures.keys(),
+    )
+
+
+new_discovery = parametrize_discovery("new discovery", "discovery_result")
+
 
 def check_categories():
     """Check that every fixture file is categorized."""
@@ -110,15 +196,15 @@ def check_categories():
         + plug.args[1]
         + bulb.args[1]
         + lightstrip.args[1]
+        + plug_smart.args[1]
     )
     diff = set(SUPPORTED_DEVICES) - set(categorized_fixtures)
     if diff:
-        for file in diff:
+        for file, protocol in diff:
             print(
-                "No category for file %s, add to the corresponding set (BULBS, PLUGS, ..)"
-                % file
+                f"No category for file {file} protocol {protocol}, add to the corresponding set (BULBS, PLUGS, ..)"
             )
-        raise Exception("Missing category for %s" % diff)
+        raise Exception(f"Missing category for {diff}")
 
 
 check_categories()
@@ -134,27 +220,32 @@ async def handle_turn_on(dev, turn_on):
         await dev.turn_off()
 
 
-def device_for_file(model):
-    for d in STRIPS:
-        if d in model:
-            return SmartStrip
+def device_for_file(model, protocol):
+    if protocol == "SMART":
+        for d in PLUGS_SMART:
+            if d in model:
+                return TapoPlug
+    else:
+        for d in STRIPS:
+            if d in model:
+                return SmartStrip
 
-    for d in PLUGS:
-        if d in model:
-            return SmartPlug
+        for d in PLUGS:
+            if d in model:
+                return SmartPlug
 
-    # Light strips are recognized also as bulbs, so this has to go first
-    for d in LIGHT_STRIPS:
-        if d in model:
-            return SmartLightStrip
+        # Light strips are recognized also as bulbs, so this has to go first
+        for d in LIGHT_STRIPS:
+            if d in model:
+                return SmartLightStrip
 
-    for d in BULBS:
-        if d in model:
-            return SmartBulb
+        for d in BULBS:
+            if d in model:
+                return SmartBulb
 
-    for d in DIMMERS:
-        if d in model:
-            return SmartDimmer
+        for d in DIMMERS:
+            if d in model:
+                return SmartDimmer
 
     raise Exception("Unable to find type for %s", model)
 
@@ -170,11 +261,14 @@ async def _discover_update_and_close(ip):
     return await _update_and_close(d)
 
 
-async def get_device_for_file(file):
+async def get_device_for_file(file, protocol):
     # if the wanted file is not an absolute path, prepend the fixtures directory
     p = Path(file)
     if not p.is_absolute():
-        p = Path(__file__).parent / "fixtures" / file
+        folder = Path(__file__).parent / "fixtures"
+        if protocol == "SMART":
+            folder = folder / "smart"
+        p = folder / file
 
     def load_file():
         with open(p) as f:
@@ -184,8 +278,12 @@ async def get_device_for_file(file):
     sysinfo = await loop.run_in_executor(None, load_file)
 
     model = basename(file)
-    d = device_for_file(model)(host="127.0.0.123")
-    d.protocol = FakeTransportProtocol(sysinfo)
+    d = device_for_file(model, protocol)(host="127.0.0.123")
+    if protocol == "SMART":
+        d.protocol = FakeSmartProtocol(sysinfo)
+        d.credentials = Credentials("", "")
+    else:
+        d.protocol = FakeTransportProtocol(sysinfo)
     await _update_and_close(d)
     return d
 
@@ -197,7 +295,7 @@ async def dev(request):
     Provides a device (given --ip) or parametrized fixture for the supported devices.
     The initial update is called automatically before returning the device.
     """
-    file = request.param
+    file, protocol = request.param
 
     ip = request.config.getoption("--ip")
     if ip:
@@ -210,19 +308,62 @@ async def dev(request):
             pytest.skip(f"skipping file {file}")
         return d if d else await _discover_update_and_close(ip)
 
-    return await get_device_for_file(file)
+    return await get_device_for_file(file, protocol)
 
 
-@pytest.fixture(params=SUPPORTED_DEVICES, scope="session")
+@pytest.fixture
+def discovery_mock(discovery_data, mocker):
+    @dataclass
+    class _DiscoveryMock:
+        ip: str
+        default_port: int
+        discovery_data: dict
+        port_override: Optional[int] = None
+
+    if "result" in discovery_data:
+        datagram = (
+            b"\x02\x00\x00\x01\x01[\x00\x00\x00\x00\x00\x00W\xcev\xf8"
+            + json_dumps(discovery_data).encode()
+        )
+        dm = _DiscoveryMock("127.0.0.123", 20002, discovery_data)
+    else:
+        datagram = TPLinkSmartHomeProtocol.encrypt(json_dumps(discovery_data))[4:]
+        dm = _DiscoveryMock("127.0.0.123", 9999, discovery_data)
+
+    def mock_discover(self):
+        port = (
+            dm.port_override
+            if dm.port_override and dm.default_port != 20002
+            else dm.default_port
+        )
+        self.datagram_received(
+            datagram,
+            (dm.ip, port),
+        )
+
+    mocker.patch("kasa.discover._DiscoverProtocol.do_discover", mock_discover)
+    mocker.patch(
+        "socket.getaddrinfo",
+        side_effect=lambda *_, **__: [(None, None, None, None, (dm.ip, 0))],
+    )
+    yield dm
+
+
+@pytest.fixture(params=FIXTURE_DATA.values(), ids=FIXTURE_DATA.keys(), scope="session")
 def discovery_data(request):
     """Return raw discovery file contents as JSON. Used for discovery tests."""
-    file = request.param
-    p = Path(file)
-    if not p.is_absolute():
-        p = Path(__file__).parent / "fixtures" / file
+    fixture_data = request.param
+    if "discovery_result" in fixture_data:
+        return {"result": fixture_data["discovery_result"]}
+    else:
+        return {"system": {"get_sysinfo": fixture_data["system"]["get_sysinfo"]}}
 
-    with open(p) as f:
-        return json.load(f)
+
+@pytest.fixture(params=FIXTURE_DATA.values(), ids=FIXTURE_DATA.keys(), scope="session")
+def all_fixture_data(request):
+    """Return raw fixture file contents as JSON. Used for discovery tests."""
+    fixture_data = request.param
+    return fixture_data
 
 
 def pytest_addoption(parser):

kasa/tests/fixtures/smart/P110_1.0_1.3.0.json

@@ -0,0 +1,180 @@
+{
+    "component_nego": {
+        "component_list": [
+            {
+                "id": "device",
+                "ver_code": 2
+            },
+            {
+                "id": "firmware",
+                "ver_code": 2
+            },
+            {
+                "id": "quick_setup",
+                "ver_code": 3
+            },
+            {
+                "id": "time",
+                "ver_code": 1
+            },
+            {
+                "id": "wireless",
+                "ver_code": 1
+            },
+            {
+                "id": "schedule",
+                "ver_code": 2
+            },
+            {
+                "id": "countdown",
+                "ver_code": 2
+            },
+            {
+                "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
+            },
+            {
+                "id": "iot_cloud",
+                "ver_code": 1
+            },
+            {
+                "id": "device_local_time",
+                "ver_code": 1
+            },
+            {
+                "id": "default_states",
+                "ver_code": 1
+            },
+            {
+                "id": "auto_off",
+                "ver_code": 2
+            },
+            {
+                "id": "localSmart",
+                "ver_code": 1
+            },
+            {
+                "id": "energy_monitoring",
+                "ver_code": 2
+            },
+            {
+                "id": "power_protection",
+                "ver_code": 1
+            },
+            {
+                "id": "current_protection",
+                "ver_code": 1
+            }
+        ]
+    },
+    "discovery_result": {
+        "device_id": "00000000000000000000000000000000",
+        "device_model": "P110(UK)",
+        "device_type": "SMART.TAPOPLUG",
+        "factory_default": false,
+        "ip": "127.0.0.123",
+        "is_support_iot_cloud": true,
+        "mac": "00-00-00-00-00-00",
+        "mgt_encrypt_schm": {
+            "encrypt_type": "KLAP",
+            "http_port": 80,
+            "is_support_https": false,
+            "lv": 2
+        },
+        "obd_src": "tplink",
+        "owner": "00000000000000000000000000000000"
+    },
+    "get_current_power": {
+        "current_power": 0
+    },
+    "get_device_info": {
+        "auto_off_remain_time": 0,
+        "auto_off_status": "off",
+        "avatar": "plug",
+        "default_states": {
+            "state": {},
+            "type": "last_states"
+        },
+        "device_id": "0000000000000000000000000000000000000000",
+        "device_on": true,
+        "fw_id": "00000000000000000000000000000000",
+        "fw_ver": "1.3.0 Build 230905 Rel.152200",
+        "has_set_location_info": true,
+        "hw_id": "00000000000000000000000000000000",
+        "hw_ver": "1.0",
+        "ip": "127.0.0.123",
+        "lang": "en_US",
+        "latitude": 0,
+        "longitude": 0,
+        "mac": "00-00-00-00-00-00",
+        "model": "P110",
+        "nickname": "VGFwaSBTbWFydCBQbHVnIDE=",
+        "oem_id": "00000000000000000000000000000000",
+        "on_time": 119335,
+        "overcurrent_status": "normal",
+        "overheated": false,
+        "power_protection_status": "normal",
+        "region": "Europe/London",
+        "rssi": -57,
+        "signal_level": 2,
+        "specs": "",
+        "ssid": "IyNNQVNLRUROQU1FIyM=",
+        "time_diff": 0,
+        "type": "SMART.TAPOPLUG"
+    },
+    "get_device_time": {
+        "region": "Europe/London",
+        "time_diff": 0,
+        "timestamp": 1701370224
+    },
+    "get_device_usage": {
+        "power_usage": {
+            "past30": 75,
+            "past7": 69,
+            "today": 0
+        },
+        "saved_power": {
+            "past30": 2029,
+            "past7": 1964,
+            "today": 1130
+        },
+        "time_usage": {
+            "past30": 2104,
+            "past7": 2033,
+            "today": 1130
+        }
+    },
+    "get_energy_usage": {
+        "current_power": 0,
+        "electricity_charge": [
+            0,
+            0,
+            0
+        ],
+        "local_time": "2023-11-30 18:50:24",
+        "month_energy": 75,
+        "month_runtime": 2104,
+        "today_energy": 0,
+        "today_runtime": 1130
+    }
+}

kasa/tests/newfakes.py

@@ -1,6 +1,7 @@
 import copy
 import logging
 import re
+from json import loads as json_loads
 
 from voluptuous import (
     REMOVE_EXTRA,
@@ -13,7 +14,8 @@ from voluptuous import (
     Schema,
 )
 
-from ..protocol import TPLinkSmartHomeProtocol
+from ..protocol import BaseTransport, TPLinkSmartHomeProtocol
+from ..smartprotocol import SmartProtocol
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -285,6 +287,41 @@ TIME_MODULE = {
 }
 
 
+class FakeSmartProtocol(SmartProtocol):
+    def __init__(self, info):
+        super().__init__("127.0.0.123", transport=FakeSmartTransport(info))
+
+
+class FakeSmartTransport(BaseTransport):
+    def __init__(self, info):
+        self.info = info
+
+    @property
+    def needs_handshake(self) -> bool:
+        return False
+
+    @property
+    def needs_login(self) -> bool:
+        return False
+
+    async def login(self, request: str) -> None:
+        pass
+
+    async def handshake(self) -> None:
+        pass
+
+    async def send(self, request: str):
+        request_dict = json_loads(request)
+        method = request_dict["method"]
+        if method == "component_nego" or method[:4] == "get_":
+            return self.info[method]
+        elif method[:4] == "set_":
+            _LOGGER.debug("Call %s not implemented, doing nothing", method)
+
+    async def close(self) -> None:
+        pass
+
+
 class FakeTransportProtocol(TPLinkSmartHomeProtocol):
     def __init__(self, info):
         self.discovery_data = info

kasa/tests/test_cli.py

@@ -6,12 +6,15 @@ from asyncclick.testing import CliRunner
 
 from kasa import SmartDevice, TPLinkSmartHomeProtocol
 from kasa.cli import alias, brightness, cli, emeter, raw_command, state, sysinfo, toggle
+from kasa.device_factory import DEVICE_TYPE_TO_CLASS
 from kasa.discover import Discover
+from kasa.smartprotocol import SmartProtocol
 
-from .conftest import handle_turn_on, turn_on
-from .newfakes import FakeTransportProtocol
+from .conftest import device_iot, handle_turn_on, new_discovery, turn_on
+from .newfakes import FakeSmartProtocol, FakeTransportProtocol
 
 
+@device_iot
 async def test_sysinfo(dev):
     runner = CliRunner()
     res = await runner.invoke(sysinfo, obj=dev)
@@ -19,6 +22,7 @@ async def test_sysinfo(dev):
     assert dev.alias in res.output
 
 
+@device_iot
 @turn_on
 async def test_state(dev, turn_on):
     await handle_turn_on(dev, turn_on)
@@ -32,6 +36,7 @@ async def test_state(dev, turn_on):
         assert "Device state: False" in res.output
 
 
+@device_iot
 @turn_on
 async def test_toggle(dev, turn_on, mocker):
     await handle_turn_on(dev, turn_on)
@@ -44,6 +49,7 @@ async def test_toggle(dev, turn_on, mocker):
         assert dev.is_on
 
 
+@device_iot
 async def test_alias(dev):
     runner = CliRunner()
 
@@ -62,6 +68,7 @@ async def test_alias(dev):
     await dev.set_alias(old_alias)
 
 
+@device_iot
 async def test_raw_command(dev):
     runner = CliRunner()
     res = await runner.invoke(raw_command, ["system", "get_sysinfo"], obj=dev)
@@ -74,6 +81,7 @@ async def test_raw_command(dev):
     assert "Usage" in res.output
 
 
+@device_iot
 async def test_emeter(dev: SmartDevice, mocker):
     runner = CliRunner()
 
@@ -99,6 +107,7 @@ async def test_emeter(dev: SmartDevice, mocker):
     daily.assert_called_with(year=1900, month=12)
 
 
+@device_iot
 async def test_brightness(dev):
     runner = CliRunner()
     res = await runner.invoke(brightness, obj=dev)
@@ -116,6 +125,7 @@ async def test_brightness(dev):
     assert "Brightness: 12" in res.output
 
 
+@device_iot
 async def test_json_output(dev: SmartDevice, mocker):
     """Test that the json output produces correct output."""
     mocker.patch("kasa.Discover.discover", return_value=[dev])
@@ -125,13 +135,9 @@ async def test_json_output(dev: SmartDevice, mocker):
     assert json.loads(res.output) == dev.internal_state
 
 
-async def test_credentials(discovery_data: dict, mocker):
+@new_discovery
+async def test_credentials(discovery_mock, mocker):
     """Test credentials are passed correctly from cli to device."""
-    # As this is testing the device constructor need to explicitly wire in
-    # the FakeTransportProtocol
-    ftp = FakeTransportProtocol(discovery_data)
-    mocker.patch.object(TPLinkSmartHomeProtocol, "query", ftp.query)
-
     # Patch state to echo username and password
     pass_dev = click.make_pass_decorator(SmartDevice)
 
@@ -143,18 +149,15 @@ async def test_credentials(discovery_data: dict, mocker):
             )
 
     mocker.patch("kasa.cli.state", new=_state)
-    cli_device_type = Discover._get_device_class(discovery_data)(
-        "any"
-    ).device_type.value
+    for subclass in DEVICE_TYPE_TO_CLASS.values():
+        mocker.patch.object(subclass, "update")
 
     runner = CliRunner()
     res = await runner.invoke(
         cli,
         [
             "--host",
-            "127.0.0.1",
-            "--type",
-            cli_device_type,
+            "127.0.0.123",
             "--username",
             "foo",
             "--password",
@@ -162,9 +165,11 @@ async def test_credentials(discovery_data: dict, mocker):
         ],
     )
     assert res.exit_code == 0
-    assert res.output == "Username:foo Password:bar\n"
+
+    assert "Username:foo Password:bar\n" in res.output
 
 
+@device_iot
 async def test_without_device_type(discovery_data: dict, dev, mocker):
     """Test connecting without the device type."""
     runner = CliRunner()

kasa/tests/test_device_factory.py

@@ -5,7 +5,9 @@ from typing import Type
 import pytest  # type: ignore # https://github.com/pytest-dev/pytest/issues/3342
 
 from kasa import (
+    Credentials,
     DeviceType,
+    Discover,
     SmartBulb,
     SmartDevice,
     SmartDeviceException,
@@ -13,8 +15,13 @@ from kasa import (
     SmartLightStrip,
     SmartPlug,
 )
-from kasa.device_factory import connect
-from kasa.klapprotocol import TPLinkKlap
+from kasa.device_factory import (
+    DEVICE_TYPE_TO_CLASS,
+    connect,
+    get_protocol_from_connection_name,
+)
+from kasa.discover import DiscoveryResult
+from kasa.iotprotocol import IotProtocol
 from kasa.protocol import TPLinkProtocol, TPLinkSmartHomeProtocol
 
 
@@ -22,11 +29,15 @@ from kasa.protocol import TPLinkProtocol, TPLinkSmartHomeProtocol
 async def test_connect(discovery_data: dict, mocker, custom_port):
     """Make sure that connect returns an initialized SmartDevice instance."""
     host = "127.0.0.1"
-    mocker.patch("kasa.TPLinkSmartHomeProtocol.query", return_value=discovery_data)
 
-    dev = await connect(host, port=custom_port)
-    assert issubclass(dev.__class__, SmartDevice)
-    assert dev.port == custom_port or dev.port == 9999
+    if "result" in discovery_data:
+        with pytest.raises(SmartDeviceException):
+            dev = await connect(host, port=custom_port)
+    else:
+        mocker.patch("kasa.TPLinkSmartHomeProtocol.query", return_value=discovery_data)
+        dev = await connect(host, port=custom_port)
+        assert issubclass(dev.__class__, SmartDevice)
+        assert dev.port == custom_port or dev.port == 9999
 
 
 @pytest.mark.parametrize("custom_port", [123, None])
@@ -49,11 +60,15 @@ async def test_connect_passed_device_type(
 ):
     """Make sure that connect with a passed device type."""
     host = "127.0.0.1"
-    mocker.patch("kasa.TPLinkSmartHomeProtocol.query", return_value=discovery_data)
 
-    dev = await connect(host, port=custom_port, device_type=device_type)
-    assert isinstance(dev, klass)
-    assert dev.port == custom_port or dev.port == 9999
+    if "result" in discovery_data:
+        with pytest.raises(SmartDeviceException):
+            dev = await connect(host, port=custom_port)
+    else:
+        mocker.patch("kasa.TPLinkSmartHomeProtocol.query", return_value=discovery_data)
+        dev = await connect(host, port=custom_port, device_type=device_type)
+        assert isinstance(dev, klass)
+        assert dev.port == custom_port or dev.port == 9999
 
 
 async def test_connect_query_fails(discovery_data: dict, mocker):
@@ -70,32 +85,52 @@ async def test_connect_logs_connect_time(
 ):
     """Test that the connect time is logged when debug logging is enabled."""
     host = "127.0.0.1"
-    mocker.patch("kasa.TPLinkSmartHomeProtocol.query", return_value=discovery_data)
-    logging.getLogger("kasa").setLevel(logging.DEBUG)
-    await connect(host)
-    assert "seconds to connect" in caplog.text
+    if "result" in discovery_data:
+        with pytest.raises(SmartDeviceException):
+            await connect(host)
+    else:
+        mocker.patch("kasa.TPLinkSmartHomeProtocol.query", return_value=discovery_data)
+        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,
+    all_fixture_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)
+    if "discovery_result" in all_fixture_data:
+        discovery_info = {"result": all_fixture_data["discovery_result"]}
+        device_class = Discover._get_device_class(discovery_info)
+    else:
+        device_class = Discover._get_device_class(all_fixture_data)
 
-    dev = await connect(host, device_type=device_type, protocol_class=protocol_in)
-    assert isinstance(dev.protocol, protocol_result)
+    device_type = list(DEVICE_TYPE_TO_CLASS.keys())[
+        list(DEVICE_TYPE_TO_CLASS.values()).index(device_class)
+    ]
+    host = "127.0.0.1"
+    if "discovery_result" in all_fixture_data:
+        mocker.patch("kasa.IotProtocol.query", return_value=all_fixture_data)
+        mocker.patch("kasa.SmartProtocol.query", return_value=all_fixture_data)
+
+        dr = DiscoveryResult(**discovery_info["result"])
+        connection_name = (
+            dr.device_type.split(".")[0] + "." + dr.mgt_encrypt_schm.encrypt_type
+        )
+        protocol_class = get_protocol_from_connection_name(
+            connection_name, host
+        ).__class__
+    else:
+        mocker.patch(
+            "kasa.TPLinkSmartHomeProtocol.query", return_value=all_fixture_data
+        )
+        protocol_class = TPLinkSmartHomeProtocol
+
+    dev = await connect(
+        host,
+        device_type=device_type,
+        protocol_class=protocol_class,
+        credentials=Credentials("", ""),
+    )
+    assert isinstance(dev.protocol, protocol_class)

kasa/tests/test_discovery.py

@@ -17,6 +17,27 @@ from kasa.exceptions import AuthenticationException, UnsupportedDeviceException
 
 from .conftest import bulb, dimmer, lightstrip, plug, strip
 
+UNSUPPORTED = {
+    "result": {
+        "device_id": "xx",
+        "owner": "xx",
+        "device_type": "SMART.TAPOXMASTREE",
+        "device_model": "P110(EU)",
+        "ip": "127.0.0.1",
+        "mac": "48-22xxx",
+        "is_support_iot_cloud": True,
+        "obd_src": "tplink",
+        "factory_default": False,
+        "mgt_encrypt_schm": {
+            "is_support_https": False,
+            "encrypt_type": "AES",
+            "http_port": 80,
+            "lv": 2,
+        },
+    },
+    "error_code": 0,
+}
+
 
 @plug
 async def test_type_detection_plug(dev: SmartDevice):
@@ -62,76 +83,40 @@ async def test_type_unknown():
 
 
 @pytest.mark.parametrize("custom_port", [123, None])
-async def test_discover_single(discovery_data: dict, mocker, custom_port):
+# @pytest.mark.parametrize("discovery_mock", [("127.0.0.1",123), ("127.0.0.1",None)], indirect=True)
+async def test_discover_single(discovery_mock, custom_port, mocker):
     """Make sure that discover_single returns an initialized SmartDevice instance."""
     host = "127.0.0.1"
-    info = {"system": {"get_sysinfo": discovery_data["system"]["get_sysinfo"]}}
-    query_mock = mocker.patch("kasa.TPLinkSmartHomeProtocol.query", return_value=info)
-
-    def mock_discover(self):
-        self.datagram_received(
-            protocol.TPLinkSmartHomeProtocol.encrypt(json_dumps(info))[4:],
-            (host, custom_port or 9999),
-        )
-
-    mocker.patch.object(_DiscoverProtocol, "do_discover", mock_discover)
+    discovery_mock.ip = host
+    discovery_mock.port_override = custom_port
+    update_mock = mocker.patch.object(SmartStrip, "update")
 
     x = await Discover.discover_single(host, port=custom_port)
     assert issubclass(x.__class__, SmartDevice)
-    assert x._sys_info is not None
-    assert x.port == custom_port or x.port == 9999
-    assert (query_mock.call_count > 0) == isinstance(x, SmartStrip)
+    assert x._discovery_info is not None
+    assert x.port == custom_port or x.port == discovery_mock.default_port
+    assert (update_mock.call_count > 0) == isinstance(x, SmartStrip)
 
 
-async def test_discover_single_hostname(discovery_data: dict, mocker):
+async def test_discover_single_hostname(discovery_mock, mocker):
     """Make sure that discover_single returns an initialized SmartDevice instance."""
     host = "foobar"
     ip = "127.0.0.1"
-    info = {"system": {"get_sysinfo": discovery_data["system"]["get_sysinfo"]}}
-    query_mock = mocker.patch("kasa.TPLinkSmartHomeProtocol.query", return_value=info)
 
-    def mock_discover(self):
-        self.datagram_received(
-            protocol.TPLinkSmartHomeProtocol.encrypt(json_dumps(info))[4:],
-            (ip, 9999),
-        )
-
-    mocker.patch.object(_DiscoverProtocol, "do_discover", mock_discover)
-    mocker.patch("socket.getaddrinfo", return_value=[(None, None, None, None, (ip, 0))])
+    discovery_mock.ip = ip
+    update_mock = mocker.patch.object(SmartStrip, "update")
 
     x = await Discover.discover_single(host)
     assert issubclass(x.__class__, SmartDevice)
-    assert x._sys_info is not None
+    assert x._discovery_info is not None
     assert x.host == host
-    assert (query_mock.call_count > 0) == isinstance(x, SmartStrip)
+    assert (update_mock.call_count > 0) == isinstance(x, SmartStrip)
 
     mocker.patch("socket.getaddrinfo", side_effect=socket.gaierror())
     with pytest.raises(SmartDeviceException):
         x = await Discover.discover_single(host)
 
 
-UNSUPPORTED = {
-    "result": {
-        "device_id": "xx",
-        "owner": "xx",
-        "device_type": "SMART.TAPOXMASTREE",
-        "device_model": "P110(EU)",
-        "ip": "127.0.0.1",
-        "mac": "48-22xxx",
-        "is_support_iot_cloud": True,
-        "obd_src": "tplink",
-        "factory_default": False,
-        "mgt_encrypt_schm": {
-            "is_support_https": False,
-            "encrypt_type": "AES",
-            "http_port": 80,
-            "lv": 2,
-        },
-    },
-    "error_code": 0,
-}
-
-
 async def test_discover_single_unsupported(mocker):
     """Make sure that discover_single handles unsupported devices correctly."""
     host = "127.0.0.1"
@@ -201,14 +186,17 @@ async def test_discover_send(mocker):
 async def test_discover_datagram_received(mocker, discovery_data):
     """Verify that datagram received fills discovered_devices."""
     proto = _DiscoverProtocol()
-    info = {"system": {"get_sysinfo": discovery_data["system"]["get_sysinfo"]}}
-    mocker.patch("kasa.discover.json_loads", return_value=info)
-    mocker.patch.object(protocol.TPLinkSmartHomeProtocol, "encrypt")
+
     mocker.patch.object(protocol.TPLinkSmartHomeProtocol, "decrypt")
 
     addr = "127.0.0.1"
-    proto.datagram_received("<placeholder data>", (addr, 9999))
+    port = 20002 if "result" in discovery_data else 9999
+
+    mocker.patch("kasa.discover.json_loads", return_value=discovery_data)
+    proto.datagram_received("<placeholder data>", (addr, port))
+
     addr2 = "127.0.0.2"
+    mocker.patch("kasa.discover.json_loads", return_value=UNSUPPORTED)
     proto.datagram_received("<placeholder data>", (addr2, 20002))
 
     # Check that device in discovered_devices is initialized correctly

kasa/tests/test_klapprotocol.py

@@ -10,9 +10,14 @@ from contextlib import nullcontext as does_not_raise
 import httpx
 import pytest
 
+from ..aestransport import AesTransport
 from ..credentials import Credentials
 from ..exceptions import AuthenticationException, SmartDeviceException
-from ..klapprotocol import KlapEncryptionSession, TPLinkKlap, _sha256
+from ..iotprotocol import IotProtocol
+from ..klaptransport import KlapEncryptionSession, KlapTransport, _sha256
+from ..smartprotocol import SmartProtocol
+
+DUMMY_QUERY = {"foobar": {"foo": "bar", "bar": "foo"}}
 
 
 class _mock_response:
@@ -21,67 +26,92 @@ class _mock_response:
         self.content = content
 
 
+@pytest.mark.parametrize("transport_class", [AesTransport, KlapTransport])
+@pytest.mark.parametrize("protocol_class", [IotProtocol, SmartProtocol])
 @pytest.mark.parametrize("retry_count", [1, 3, 5])
-async def test_protocol_retries(mocker, retry_count):
+async def test_protocol_retries(mocker, retry_count, protocol_class, transport_class):
+    host = "127.0.0.1"
     conn = mocker.patch.object(
-        TPLinkKlap, "client_post", side_effect=Exception("dummy exception")
+        transport_class, "client_post", side_effect=Exception("dummy exception")
     )
     with pytest.raises(SmartDeviceException):
-        await TPLinkKlap("127.0.0.1").query({}, retry_count=retry_count)
+        await protocol_class(host, transport=transport_class(host)).query(
+            DUMMY_QUERY, retry_count=retry_count
+        )
 
     assert conn.call_count == retry_count + 1
 
 
-async def test_protocol_no_retry_on_connection_error(mocker):
+@pytest.mark.parametrize("transport_class", [AesTransport, KlapTransport])
+@pytest.mark.parametrize("protocol_class", [IotProtocol, SmartProtocol])
+async def test_protocol_no_retry_on_connection_error(
+    mocker, protocol_class, transport_class
+):
+    host = "127.0.0.1"
     conn = mocker.patch.object(
-        TPLinkKlap,
+        transport_class,
         "client_post",
         side_effect=httpx.ConnectError("foo"),
     )
     with pytest.raises(SmartDeviceException):
-        await TPLinkKlap("127.0.0.1").query({}, retry_count=5)
+        await protocol_class(host, transport=transport_class(host)).query(
+            DUMMY_QUERY, retry_count=5
+        )
 
     assert conn.call_count == 1
 
 
-async def test_protocol_retry_recoverable_error(mocker):
+@pytest.mark.parametrize("transport_class", [AesTransport, KlapTransport])
+@pytest.mark.parametrize("protocol_class", [IotProtocol, SmartProtocol])
+async def test_protocol_retry_recoverable_error(
+    mocker, protocol_class, transport_class
+):
+    host = "127.0.0.1"
     conn = mocker.patch.object(
-        TPLinkKlap,
+        transport_class,
         "client_post",
         side_effect=httpx.CloseError("foo"),
     )
     with pytest.raises(SmartDeviceException):
-        await TPLinkKlap("127.0.0.1").query({}, retry_count=5)
+        await protocol_class(host, transport=transport_class(host)).query(
+            DUMMY_QUERY, retry_count=5
+        )
 
     assert conn.call_count == 6
 
 
+@pytest.mark.parametrize("transport_class", [AesTransport, KlapTransport])
+@pytest.mark.parametrize("protocol_class", [IotProtocol, SmartProtocol])
 @pytest.mark.parametrize("retry_count", [1, 3, 5])
-async def test_protocol_reconnect(mocker, retry_count):
+async def test_protocol_reconnect(mocker, retry_count, protocol_class, transport_class):
+    host = "127.0.0.1"
     remaining = retry_count
+    mock_response = {"result": {"great": "success"}}
 
     def _fail_one_less_than_retry_count(*_, **__):
-        nonlocal remaining, encryption_session
+        nonlocal remaining
         remaining -= 1
         if remaining:
             raise Exception("Simulated post failure")
-        # Do the encrypt just before returning the value so the incrementing sequence number is correct
-        encrypted, seq = encryption_session.encrypt('{"great":"success"}')
-        return 200, encrypted
 
-    seed = secrets.token_bytes(16)
-    auth_hash = TPLinkKlap.generate_auth_hash(Credentials("foo", "bar"))
-    encryption_session = KlapEncryptionSession(seed, seed, auth_hash)
-    protocol = TPLinkKlap("127.0.0.1")
-    protocol.handshake_done = True
-    protocol.session_expire_at = time.time() + 86400
-    protocol.encryption_session = encryption_session
+        return mock_response
+
     mocker.patch.object(
-        TPLinkKlap, "client_post", side_effect=_fail_one_less_than_retry_count
+        transport_class, "needs_handshake", property(lambda self: False)
+    )
+    mocker.patch.object(transport_class, "needs_login", property(lambda self: False))
+
+    send_mock = mocker.patch.object(
+        transport_class,
+        "send",
+        side_effect=_fail_one_less_than_retry_count,
     )
 
-    response = await protocol.query({}, retry_count=retry_count)
-    assert response == {"great": "success"}
+    response = await protocol_class(host, transport=transport_class(host)).query(
+        DUMMY_QUERY, retry_count=retry_count
+    )
+    assert "result" in response or "great" in response
+    assert send_mock.call_count == retry_count
 
 
 @pytest.mark.parametrize("log_level", [logging.WARNING, logging.DEBUG])
@@ -96,14 +126,14 @@ async def test_protocol_logging(mocker, caplog, log_level):
         return 200, encrypted
 
     seed = secrets.token_bytes(16)
-    auth_hash = TPLinkKlap.generate_auth_hash(Credentials("foo", "bar"))
+    auth_hash = KlapTransport.generate_auth_hash(Credentials("foo", "bar"))
     encryption_session = KlapEncryptionSession(seed, seed, auth_hash)
-    protocol = TPLinkKlap("127.0.0.1")
+    protocol = IotProtocol("127.0.0.1")
 
-    protocol.handshake_done = True
-    protocol.session_expire_at = time.time() + 86400
-    protocol.encryption_session = encryption_session
-    mocker.patch.object(TPLinkKlap, "client_post", side_effect=_return_encrypted)
+    protocol._transport._handshake_done = True
+    protocol._transport._session_expire_at = time.time() + 86400
+    protocol._transport._encryption_session = encryption_session
+    mocker.patch.object(KlapTransport, "client_post", side_effect=_return_encrypted)
 
     response = await protocol.query({})
     assert response == {"great": "success"}
@@ -117,7 +147,7 @@ def test_encrypt():
     d = json.dumps({"foo": 1, "bar": 2})
 
     seed = secrets.token_bytes(16)
-    auth_hash = TPLinkKlap.generate_auth_hash(Credentials("foo", "bar"))
+    auth_hash = KlapTransport.generate_auth_hash(Credentials("foo", "bar"))
     encryption_session = KlapEncryptionSession(seed, seed, auth_hash)
 
     encrypted, seq = encryption_session.encrypt(d)
@@ -129,7 +159,7 @@ def test_encrypt_unicode():
     d = "{'snowman': '\u2603'}"
 
     seed = secrets.token_bytes(16)
-    auth_hash = TPLinkKlap.generate_auth_hash(Credentials("foo", "bar"))
+    auth_hash = KlapTransport.generate_auth_hash(Credentials("foo", "bar"))
     encryption_session = KlapEncryptionSession(seed, seed, auth_hash)
 
     encrypted, seq = encryption_session.encrypt(d)
@@ -145,7 +175,10 @@ def test_encrypt_unicode():
         (Credentials("foo", "bar"), does_not_raise()),
         (Credentials("", ""), does_not_raise()),
         (
-            Credentials(TPLinkKlap.KASA_SETUP_EMAIL, TPLinkKlap.KASA_SETUP_PASSWORD),
+            Credentials(
+                KlapTransport.KASA_SETUP_EMAIL,
+                KlapTransport.KASA_SETUP_PASSWORD,
+            ),
             does_not_raise(),
         ),
         (
@@ -167,21 +200,21 @@ async def test_handshake1(mocker, device_credentials, expectation):
     client_seed = None
     server_seed = secrets.token_bytes(16)
     client_credentials = Credentials("foo", "bar")
-    device_auth_hash = TPLinkKlap.generate_auth_hash(device_credentials)
+    device_auth_hash = KlapTransport.generate_auth_hash(device_credentials)
 
     mocker.patch.object(
         httpx.AsyncClient, "post", side_effect=_return_handshake1_response
     )
 
-    protocol = TPLinkKlap("127.0.0.1", credentials=client_credentials)
+    protocol = IotProtocol("127.0.0.1", credentials=client_credentials)
 
-    protocol.http_client = httpx.AsyncClient()
+    protocol._transport.http_client = httpx.AsyncClient()
     with expectation:
         (
             local_seed,
             device_remote_seed,
             auth_hash,
-        ) = await protocol.perform_handshake1()
+        ) = await protocol._transport.perform_handshake1()
 
         assert local_seed == client_seed
         assert device_remote_seed == server_seed
@@ -204,23 +237,23 @@ async def test_handshake(mocker):
     client_seed = None
     server_seed = secrets.token_bytes(16)
     client_credentials = Credentials("foo", "bar")
-    device_auth_hash = TPLinkKlap.generate_auth_hash(client_credentials)
+    device_auth_hash = KlapTransport.generate_auth_hash(client_credentials)
 
     mocker.patch.object(
         httpx.AsyncClient, "post", side_effect=_return_handshake_response
     )
 
-    protocol = TPLinkKlap("127.0.0.1", credentials=client_credentials)
-    protocol.http_client = httpx.AsyncClient()
+    protocol = IotProtocol("127.0.0.1", credentials=client_credentials)
+    protocol._transport.http_client = httpx.AsyncClient()
 
     response_status = 200
-    await protocol.perform_handshake()
-    assert protocol.handshake_done is True
+    await protocol._transport.perform_handshake()
+    assert protocol._transport._handshake_done is True
 
     response_status = 403
     with pytest.raises(AuthenticationException):
-        await protocol.perform_handshake()
-    assert protocol.handshake_done is False
+        await protocol._transport.perform_handshake()
+    assert protocol._transport._handshake_done is False
     await protocol.close()
 
 
@@ -237,9 +270,9 @@ async def test_query(mocker):
             return _mock_response(200, b"")
         elif url == "http://127.0.0.1/app/request":
             encryption_session = KlapEncryptionSession(
-                protocol.encryption_session.local_seed,
-                protocol.encryption_session.remote_seed,
-                protocol.encryption_session.user_hash,
+                protocol._transport._encryption_session.local_seed,
+                protocol._transport._encryption_session.remote_seed,
+                protocol._transport._encryption_session.user_hash,
             )
             seq = params.get("seq")
             encryption_session._seq = seq - 1
@@ -252,11 +285,11 @@ async def test_query(mocker):
     seq = None
     server_seed = secrets.token_bytes(16)
     client_credentials = Credentials("foo", "bar")
-    device_auth_hash = TPLinkKlap.generate_auth_hash(client_credentials)
+    device_auth_hash = KlapTransport.generate_auth_hash(client_credentials)
 
     mocker.patch.object(httpx.AsyncClient, "post", side_effect=_return_response)
 
-    protocol = TPLinkKlap("127.0.0.1", credentials=client_credentials)
+    protocol = IotProtocol("127.0.0.1", credentials=client_credentials)
 
     for _ in range(10):
         resp = await protocol.query({})
@@ -296,11 +329,11 @@ async def test_authentication_failures(mocker, response_status, expectation):
 
     server_seed = secrets.token_bytes(16)
     client_credentials = Credentials("foo", "bar")
-    device_auth_hash = TPLinkKlap.generate_auth_hash(client_credentials)
+    device_auth_hash = KlapTransport.generate_auth_hash(client_credentials)
 
     mocker.patch.object(httpx.AsyncClient, "post", side_effect=_return_response)
 
-    protocol = TPLinkKlap("127.0.0.1", credentials=client_credentials)
+    protocol = IotProtocol("127.0.0.1", credentials=client_credentials)
 
     with expectation:
         await protocol.query({})

kasa/tests/test_plug.py

@@ -1,6 +1,6 @@
 from kasa import DeviceType
 
-from .conftest import plug
+from .conftest import plug, plug_smart
 from .newfakes import PLUG_SCHEMA
 
 
@@ -28,3 +28,14 @@ async def test_led(dev):
     assert dev.led
 
     await dev.set_led(original)
+
+
+@plug_smart
+async def test_plug_device_info(dev):
+    assert dev._info is not None
+    # PLUG_SCHEMA(dev.sys_info)
+
+    assert dev.model is not None
+
+    assert dev.device_type == DeviceType.Plug or dev.device_type == DeviceType.Strip
+    # assert dev.is_plug or dev.is_strip

kasa/tests/test_readme_examples.py

@@ -9,7 +9,7 @@ from kasa.tests.conftest import get_device_for_file
 
 def test_bulb_examples(mocker):
     """Use KL130 (bulb with all features) to test the doctests."""
-    p = asyncio.run(get_device_for_file("KL130(US)_1.0_1.8.11.json"))
+    p = asyncio.run(get_device_for_file("KL130(US)_1.0_1.8.11.json", "IOT"))
     mocker.patch("kasa.smartbulb.SmartBulb", return_value=p)
     mocker.patch("kasa.smartbulb.SmartBulb.update")
     res = xdoctest.doctest_module("kasa.smartbulb", "all")
@@ -18,7 +18,7 @@ def test_bulb_examples(mocker):
 
 def test_smartdevice_examples(mocker):
     """Use HS110 for emeter examples."""
-    p = asyncio.run(get_device_for_file("HS110(EU)_1.0_1.2.5.json"))
+    p = asyncio.run(get_device_for_file("HS110(EU)_1.0_1.2.5.json", "IOT"))
     mocker.patch("kasa.smartdevice.SmartDevice", return_value=p)
     mocker.patch("kasa.smartdevice.SmartDevice.update")
     res = xdoctest.doctest_module("kasa.smartdevice", "all")
@@ -27,7 +27,7 @@ def test_smartdevice_examples(mocker):
 
 def test_plug_examples(mocker):
     """Test plug examples."""
-    p = asyncio.run(get_device_for_file("HS110(EU)_1.0_1.2.5.json"))
+    p = asyncio.run(get_device_for_file("HS110(EU)_1.0_1.2.5.json", "IOT"))
     mocker.patch("kasa.smartplug.SmartPlug", return_value=p)
     mocker.patch("kasa.smartplug.SmartPlug.update")
     res = xdoctest.doctest_module("kasa.smartplug", "all")
@@ -36,7 +36,7 @@ def test_plug_examples(mocker):
 
 def test_strip_examples(mocker):
     """Test strip examples."""
-    p = asyncio.run(get_device_for_file("KP303(UK)_1.0_1.0.3.json"))
+    p = asyncio.run(get_device_for_file("KP303(UK)_1.0_1.0.3.json", "IOT"))
     mocker.patch("kasa.smartstrip.SmartStrip", return_value=p)
     mocker.patch("kasa.smartstrip.SmartStrip.update")
     res = xdoctest.doctest_module("kasa.smartstrip", "all")
@@ -45,7 +45,7 @@ def test_strip_examples(mocker):
 
 def test_dimmer_examples(mocker):
     """Test dimmer examples."""
-    p = asyncio.run(get_device_for_file("HS220(US)_1.0_1.5.7.json"))
+    p = asyncio.run(get_device_for_file("HS220(US)_1.0_1.5.7.json", "IOT"))
     mocker.patch("kasa.smartdimmer.SmartDimmer", return_value=p)
     mocker.patch("kasa.smartdimmer.SmartDimmer.update")
     res = xdoctest.doctest_module("kasa.smartdimmer", "all")
@@ -54,7 +54,7 @@ def test_dimmer_examples(mocker):
 
 def test_lightstrip_examples(mocker):
     """Test lightstrip examples."""
-    p = asyncio.run(get_device_for_file("KL430(US)_1.0_1.0.10.json"))
+    p = asyncio.run(get_device_for_file("KL430(US)_1.0_1.0.10.json", "IOT"))
     mocker.patch("kasa.smartlightstrip.SmartLightStrip", return_value=p)
     mocker.patch("kasa.smartlightstrip.SmartLightStrip.update")
     res = xdoctest.doctest_module("kasa.smartlightstrip", "all")
@@ -63,7 +63,7 @@ def test_lightstrip_examples(mocker):
 
 def test_discovery_examples(mocker):
     """Test discovery examples."""
-    p = asyncio.run(get_device_for_file("KP303(UK)_1.0_1.0.3.json"))
+    p = asyncio.run(get_device_for_file("KP303(UK)_1.0_1.0.3.json", "IOT"))
 
     mocker.patch("kasa.discover.Discover.discover", return_value=[p])
     res = xdoctest.doctest_module("kasa.discover", "all")

kasa/tests/test_smartdevice.py

@@ -8,7 +8,7 @@ import kasa
 from kasa import Credentials, SmartDevice, SmartDeviceException
 from kasa.smartdevice import DeviceType
 
-from .conftest import handle_turn_on, has_emeter, no_emeter, turn_on
+from .conftest import device_iot, handle_turn_on, has_emeter, no_emeter, turn_on
 from .newfakes import PLUG_SCHEMA, TZ_SCHEMA, FakeTransportProtocol
 
 # List of all SmartXXX classes including the SmartDevice base class
@@ -22,11 +22,13 @@ smart_device_classes = [
 ]
 
 
+@device_iot
 async def test_state_info(dev):
     assert isinstance(dev.state_information, dict)
 
 
 @pytest.mark.requires_dummy
+@device_iot
 async def test_invalid_connection(dev):
     with patch.object(
         FakeTransportProtocol, "query", side_effect=SmartDeviceException
@@ -58,12 +60,14 @@ async def test_initial_update_no_emeter(dev, mocker):
     assert spy.call_count == 2
 
 
+@device_iot
 async def test_query_helper(dev):
     with pytest.raises(SmartDeviceException):
         await dev._query_helper("test", "testcmd", {})
     # TODO check for unwrapping?
 
 
+@device_iot
 @turn_on
 async def test_state(dev, turn_on):
     await handle_turn_on(dev, turn_on)
@@ -90,6 +94,7 @@ async def test_state(dev, turn_on):
         assert dev.is_off
 
 
+@device_iot
 async def test_alias(dev):
     test_alias = "TEST1234"
     original = dev.alias
@@ -104,6 +109,7 @@ async def test_alias(dev):
     assert dev.alias == original
 
 
+@device_iot
 @turn_on
 async def test_on_since(dev, turn_on):
     await handle_turn_on(dev, turn_on)
@@ -116,30 +122,37 @@ async def test_on_since(dev, turn_on):
         assert dev.on_since is None
 
 
+@device_iot
 async def test_time(dev):
     assert isinstance(await dev.get_time(), datetime)
 
 
+@device_iot
 async def test_timezone(dev):
     TZ_SCHEMA(await dev.get_timezone())
 
 
+@device_iot
 async def test_hw_info(dev):
     PLUG_SCHEMA(dev.hw_info)
 
 
+@device_iot
 async def test_location(dev):
     PLUG_SCHEMA(dev.location)
 
 
+@device_iot
 async def test_rssi(dev):
     PLUG_SCHEMA({"rssi": dev.rssi})  # wrapping for vol
 
 
+@device_iot
 async def test_mac(dev):
     PLUG_SCHEMA({"mac": dev.mac})  # wrapping for val
 
 
+@device_iot
 async def test_representation(dev):
     import re
 
@@ -147,6 +160,7 @@ async def test_representation(dev):
     assert pattern.match(str(dev))
 
 
+@device_iot
 async def test_childrens(dev):
     """Make sure that children property is exposed by every device."""
     if dev.is_strip:
@@ -155,6 +169,7 @@ async def test_childrens(dev):
         assert len(dev.children) == 0
 
 
+@device_iot
 async def test_children(dev):
     """Make sure that children property is exposed by every device."""
     if dev.is_strip:
@@ -165,11 +180,13 @@ async def test_children(dev):
         assert dev.has_children is False
 
 
+@device_iot
 async def test_internal_state(dev):
     """Make sure the internal state returns the last update results."""
     assert dev.internal_state == dev._last_update
 
 
+@device_iot
 async def test_features(dev):
     """Make sure features is always accessible."""
     sysinfo = dev._last_update["system"]["get_sysinfo"]
@@ -179,11 +196,13 @@ async def test_features(dev):
         assert dev.features == set()
 
 
+@device_iot
 async def test_max_device_response_size(dev):
     """Make sure every device return has a set max response size."""
     assert dev.max_device_response_size > 0
 
 
+@device_iot
 async def test_estimated_response_sizes(dev):
     """Make sure every module has an estimated response size set."""
     for mod in dev.modules.values():
@@ -202,6 +221,7 @@ def test_device_class_ctors(device_class):
     assert dev.credentials == credentials
 
 
+@device_iot
 async def test_modules_preserved(dev: SmartDevice):
     """Make modules that are not being updated are preserved between updates."""
     dev._last_update["some_module_not_being_updated"] = "should_be_kept"
@@ -237,6 +257,7 @@ async def test_create_thin_wrapper():
     )
 
 
+@device_iot
 async def test_modules_not_supported(dev: SmartDevice):
     """Test that unsupported modules do not break the device."""
     for module in dev.modules.values():