python-kasa/tests/transports/test_klaptransport.py

from __future__ import annotations

import json
import logging
import re
import secrets
import time
from collections.abc import Callable
from contextlib import AbstractContextManager
from contextlib import nullcontext as does_not_raise

import aiohttp
import pytest
from cryptography.hazmat.primitives.ciphers import Cipher, modes
from pytest_mock import MockerFixture
from yarl import URL

from kasa.credentials import DEFAULT_CREDENTIALS, Credentials, get_default_credentials
from kasa.deviceconfig import (
    DeviceConfig,
    DeviceConnectionParameters,
    DeviceEncryptionType,
    DeviceFamily,
)
from kasa.exceptions import (
    AuthenticationError,
    KasaException,
    TimeoutError,
    _ConnectionError,
    _RetryableError,
)
from kasa.httpclient import HttpClient
from kasa.iot import IotBulb, IotDevice, IotPlug, IotWallSwitch
from kasa.protocols import IotProtocol, SmartProtocol
from kasa.transports.aestransport import AesTransport
from kasa.transports.klaptransport import (
    PACK_SIGNED_LONG,
    KlapEncryptionSession,
    KlapTransport,
    KlapTransportV2,
    _sha256,
)

from ..device_fixtures import device_iot, parametrize

DUMMY_QUERY = {"foobar": {"foo": "bar", "bar": "foo"}}

KLAP_CONNECTION = DeviceConnectionParameters(
    DeviceFamily.IotSmartPlugSwitch, DeviceEncryptionType.Klap, login_version=2
)

ks200m_iot = parametrize("ks200m", model_filter={"KS200M"}, protocol_filter={"IOT"})

# Transport tests are not designed for real devices
pytestmark = [pytest.mark.requires_dummy]


class _mock_response:
    def __init__(self, status, content: bytes):
        self.status = status
        self.content = content

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_t, exc_v, exc_tb):
        pass

    async def read(self):
        return self.content


@pytest.mark.parametrize(
    ("error", "retry_expectation"),
    [
        (Exception("dummy exception"), False),
        (aiohttp.ServerTimeoutError("dummy exception"), True),
        (aiohttp.ServerDisconnectedError("dummy exception"), True),
        (aiohttp.ClientOSError("dummy exception"), True),
    ],
    ids=("Exception", "ServerTimeoutError", "ServerDisconnectedError", "ClientOSError"),
)
@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_via_client_session(
    mocker: MockerFixture,
    retry_count: int,
    protocol_class: type,
    transport_class: type,
    error: Exception,
    retry_expectation: bool,
) -> None:
    host = "127.0.0.1"
    conn = mocker.patch.object(aiohttp.ClientSession, "post", side_effect=error)

    config = DeviceConfig(host)
    with pytest.raises(KasaException):
        await protocol_class(transport=transport_class(config=config)).query(
            DUMMY_QUERY, retry_count=retry_count
        )

    expected_count = retry_count + 1 if retry_expectation else 1
    assert conn.call_count == expected_count


@pytest.mark.parametrize(
    ("error", "retry_expectation"),
    [
        (KasaException("dummy exception"), False),
        (_RetryableError("dummy exception"), True),
        (TimeoutError("dummy exception"), True),
    ],
    ids=("KasaException", "_RetryableError", "TimeoutError"),
)
@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_via_httpclient(
    mocker: MockerFixture,
    retry_count: int,
    protocol_class: type,
    transport_class: type,
    error: Exception,
    retry_expectation: bool,
) -> None:
    host = "127.0.0.1"
    conn = mocker.patch.object(HttpClient, "post", side_effect=error)
    mocker.patch.object(protocol_class, "BACKOFF_SECONDS_AFTER_TIMEOUT", 0)

    config = DeviceConfig(host)
    with pytest.raises(KasaException):
        await protocol_class(transport=transport_class(config=config)).query(
            DUMMY_QUERY, retry_count=retry_count
        )

    expected_count = retry_count + 1 if retry_expectation else 1
    assert conn.call_count == expected_count


@pytest.mark.parametrize("transport_class", [AesTransport, KlapTransport])
@pytest.mark.parametrize("protocol_class", [IotProtocol, SmartProtocol])
async def test_protocol_no_retry_on_connection_error(
    mocker: MockerFixture, protocol_class: type, transport_class: type
) -> None:
    host = "127.0.0.1"
    conn = mocker.patch.object(
        aiohttp.ClientSession,
        "post",
        side_effect=AuthenticationError("foo"),
    )
    mocker.patch.object(protocol_class, "BACKOFF_SECONDS_AFTER_TIMEOUT", 0)
    config = DeviceConfig(host)
    with pytest.raises(KasaException):
        await protocol_class(transport=transport_class(config=config)).query(
            DUMMY_QUERY, retry_count=5
        )

    assert conn.call_count == 1


@pytest.mark.parametrize("transport_class", [AesTransport, KlapTransport])
@pytest.mark.parametrize("protocol_class", [IotProtocol, SmartProtocol])
async def test_protocol_retry_recoverable_error(
    mocker: MockerFixture, protocol_class: type, transport_class: type
) -> None:
    host = "127.0.0.1"
    conn = mocker.patch.object(
        aiohttp.ClientSession,
        "post",
        side_effect=aiohttp.ClientOSError("foo"),
    )

    config = DeviceConfig(host)
    with pytest.raises(KasaException):
        await protocol_class(transport=transport_class(config=config)).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: MockerFixture, retry_count: int, protocol_class: type, transport_class: type
) -> None:
    host = "127.0.0.1"
    remaining = retry_count
    mock_response = {"result": {"great": "success"}, "error_code": 0}

    def _fail_one_less_than_retry_count(*_, **__):
        nonlocal remaining
        remaining -= 1
        if remaining:
            raise _ConnectionError("Simulated connection failure")

        return mock_response

    mocker.patch.object(transport_class, "perform_handshake")
    if hasattr(transport_class, "perform_login"):
        mocker.patch.object(transport_class, "perform_login")

    send_mock = mocker.patch.object(
        transport_class,
        "send",
        side_effect=_fail_one_less_than_retry_count,
    )

    config = DeviceConfig(host)
    response = await protocol_class(transport=transport_class(config=config)).query(
        DUMMY_QUERY, retry_count=retry_count
    )
    assert "result" in response or "foobar" in response
    assert send_mock.call_count == retry_count


@pytest.mark.parametrize("log_level", [logging.WARNING, logging.DEBUG])
@pytest.mark.xdist_group(name="caplog")
async def test_protocol_logging(
    mocker: MockerFixture, caplog: pytest.LogCaptureFixture, log_level: int
) -> None:
    caplog.set_level(log_level)
    logging.getLogger("kasa").setLevel(log_level)

    def _return_encrypted(*_, **__):
        # Devices sign the response for the sequence number of the request
        encryption_session._seq -= 1
        encrypted, seq = encryption_session.encrypt('{"great":"success"}')
        return 200, encrypted

    seed = secrets.token_bytes(16)
    auth_hash = KlapTransport.generate_auth_hash(Credentials("foo", "bar"))
    encryption_session = KlapEncryptionSession(seed, seed, auth_hash)

    config = DeviceConfig("127.0.0.1")
    transport = KlapTransport(config=config)
    protocol = IotProtocol(transport=transport)

    transport._handshake_done = True
    transport._session_expire_at = time.time() + 86400
    transport._encryption_session = encryption_session
    mocker.patch.object(HttpClient, "post", side_effect=_return_encrypted)

    response = await protocol.query({})
    assert response == {"great": "success"}
    if log_level == logging.DEBUG:
        assert "success" in caplog.text
    else:
        assert "success" not in caplog.text


def test_encrypt() -> None:
    d = json.dumps({"foo": 1, "bar": 2})

    seed = secrets.token_bytes(16)
    auth_hash = KlapTransport.generate_auth_hash(Credentials("foo", "bar"))
    encryption_session = KlapEncryptionSession(seed, seed, auth_hash)

    encrypted, seq = encryption_session.encrypt(d)

    assert d == encryption_session.decrypt(encrypted)


def test_encrypt_unicode() -> None:
    d = "{'snowman': '\u2603'}"

    seed = secrets.token_bytes(16)
    auth_hash = KlapTransport.generate_auth_hash(Credentials("foo", "bar"))
    encryption_session = KlapEncryptionSession(seed, seed, auth_hash)

    encrypted, seq = encryption_session.encrypt(d)

    decrypted = encryption_session.decrypt(encrypted)

    assert d == decrypted


async def test_transport_decrypt(mocker: MockerFixture) -> None:
    """Test transport decryption."""
    d = {"great": "success"}

    seed = secrets.token_bytes(16)
    auth_hash = KlapTransport.generate_auth_hash(Credentials("foo", "bar"))
    encryption_session = KlapEncryptionSession(seed, seed, auth_hash)

    transport = KlapTransport(config=DeviceConfig(host="127.0.0.1"))
    transport._handshake_done = True
    transport._session_expire_at = time.monotonic() + 60
    transport._encryption_session = encryption_session

    async def _return_response(url: URL, params=None, data=None, *_, **__):
        assert transport._encryption_session is not None
        encryption_session = KlapEncryptionSession(
            transport._encryption_session.local_seed,
            transport._encryption_session.remote_seed,
            transport._encryption_session.user_hash,
        )
        seq = params.get("seq")
        encryption_session._seq = seq - 1
        encrypted, seq = encryption_session.encrypt(json.dumps(d))
        seq = seq
        return 200, encrypted

    mocker.patch.object(HttpClient, "post", side_effect=_return_response)

    resp = await transport.send(json.dumps({}))
    assert d == resp


def _logged_in_transport() -> KlapTransport:
    seed = secrets.token_bytes(16)
    auth_hash = KlapTransport.generate_auth_hash(Credentials("foo", "bar"))

    transport = KlapTransport(config=DeviceConfig(host="127.0.0.1"))
    transport._handshake_done = True
    transport._session_expire_at = time.monotonic() + 60
    transport._encryption_session = KlapEncryptionSession(seed, seed, auth_hash)
    return transport


@pytest.mark.parametrize("mismatch", ["key", "seq"])
async def test_transport_signature_error(mocker: MockerFixture, mismatch: str) -> None:
    """Test that a response not signed for the request raises a kasa exception."""
    transport = _logged_in_transport()

    async def _return_response(url: URL, params=None, data=None, *_, **__):
        assert transport._encryption_session is not None
        encryption_session = KlapEncryptionSession(
            secrets.token_bytes(16)
            if mismatch == "key"
            else transport._encryption_session.local_seed,
            transport._encryption_session.remote_seed,
            transport._encryption_session.user_hash,
        )
        seq = params.get("seq")
        encryption_session._seq = seq - 1 if mismatch == "key" else seq
        encrypted, _seq = encryption_session.encrypt(json.dumps({"great": "success"}))
        return 200, encrypted

    mocker.patch.object(HttpClient, "post", side_effect=_return_response)

    with pytest.raises(
        KasaException,
        match=re.escape(
            "Device 127.0.0.1 sent a response with an invalid signature "
            "for request with seq"
        ),
    ):
        await transport.send(json.dumps({}))


async def test_transport_decrypt_error(mocker: MockerFixture) -> None:
    """Test that a signed response that fails to decrypt raises a kasa exception."""
    transport = _logged_in_transport()

    async def _return_response(url: URL, params=None, data=None, *_, **__):
        session = transport._encryption_session
        assert session is not None
        seq = PACK_SIGNED_LONG(params.get("seq"))
        # A block that decrypts to zeros, which is not valid PKCS7 padding
        encryptor = Cipher(session._aes, modes.CBC(session._iv + seq)).encryptor()
        ciphertext = encryptor.update(bytes(16)) + encryptor.finalize()
        return 200, _sha256(session._sig + seq + ciphertext) + ciphertext

    mocker.patch.object(HttpClient, "post", side_effect=_return_response)

    with pytest.raises(
        KasaException,
        match=re.escape("Error trying to decrypt device 127.0.0.1 response:"),
    ):
        await transport.send(json.dumps({}))


HTML_PAGE = b"<html><body><center>200 OK</center></body></html>"


@pytest.mark.parametrize(
    "response",
    [
        pytest.param(HTML_PAGE, id="html-page"),
        pytest.param(b"", id="empty"),
        pytest.param(secrets.token_bytes(32), id="signature-only"),
        pytest.param(secrets.token_bytes(32 + 17), id="partial-block"),
    ],
)
async def test_transport_invalid_response(
    mocker: MockerFixture, response: bytes
) -> None:
    """Test that a response that is not a KLAP message raises a kasa exception."""
    transport = _logged_in_transport()
    mocker.patch.object(HttpClient, "post", return_value=(200, response))

    with pytest.raises(
        KasaException,
        match=re.escape(
            "Device 127.0.0.1 sent a response that is not a KLAP message "
            f"({len(response)} bytes): {response[:64]!r}"
        ),
    ):
        await transport.send(json.dumps({}))


@pytest.mark.parametrize("transport_class", [KlapTransport, KlapTransportV2])
async def test_protocol_invalid_response(
    mocker: MockerFixture, transport_class: type[KlapTransport]
) -> None:
    """Test that an invalid response is not retried and resets the session."""
    server_seed = secrets.token_bytes(16)
    client_credentials = Credentials("foo", "bar")
    device_auth_hash = transport_class.generate_auth_hash(client_credentials)
    handshakes = 0
    requests = 0

    async def _return_response(url: URL, params=None, data=None, *_, **__):
        nonlocal handshakes, requests

        if url == URL("http://127.0.0.1:80/app/handshake1"):
            handshakes += 1
            seed_auth_hash = _sha256(
                data + server_seed + device_auth_hash
                if transport_class is KlapTransportV2
                else data + device_auth_hash
            )
            return _mock_response(200, server_seed + seed_auth_hash)
        elif url == URL("http://127.0.0.1:80/app/handshake2"):
            return _mock_response(200, b"")
        elif url == URL("http://127.0.0.1:80/app/request"):
            requests += 1
            if requests == 1:
                return _mock_response(200, HTML_PAGE)
            assert transport._encryption_session is not None
            encryption_session = KlapEncryptionSession(
                transport._encryption_session.local_seed,
                transport._encryption_session.remote_seed,
                transport._encryption_session.user_hash,
            )
            encryption_session._seq = params.get("seq") - 1
            encrypted, _seq = encryption_session.encrypt('{"great": "success"}')
            return _mock_response(200, encrypted)

    mocker.patch.object(aiohttp.ClientSession, "post", side_effect=_return_response)

    config = DeviceConfig("127.0.0.1", credentials=client_credentials)
    transport = transport_class(config=config)
    protocol = IotProtocol(transport=transport)

    with pytest.raises(KasaException, match="is not a KLAP message"):
        await protocol.query(DUMMY_QUERY)
    assert requests == 1
    assert handshakes == 1
    assert transport._handshake_done is False

    assert await protocol.query(DUMMY_QUERY) == {"great": "success"}
    assert requests == 2
    assert handshakes == 2


@pytest.mark.parametrize(
    ("device_credentials", "expectation"),
    [
        (Credentials("foo", "bar"), does_not_raise()),
        (Credentials(), does_not_raise()),
        (
            get_default_credentials(DEFAULT_CREDENTIALS["KASA"]),
            does_not_raise(),
        ),
        (
            Credentials("shouldfail", "shouldfail"),
            pytest.raises(AuthenticationError),
        ),
    ],
    ids=("client", "blank", "kasa_setup", "shouldfail"),
)
@pytest.mark.parametrize(
    ("transport_class", "seed_auth_hash_calc"),
    [
        pytest.param(KlapTransport, lambda c, s, a: c + a, id="KLAP"),
        pytest.param(KlapTransportV2, lambda c, s, a: c + s + a, id="KLAPV2"),
    ],
)
async def test_handshake1(
    mocker: MockerFixture,
    device_credentials: Credentials,
    expectation: AbstractContextManager,
    transport_class: type[KlapTransport],
    seed_auth_hash_calc: Callable[..., bytes],
) -> None:
    client_seed = None
    server_seed = secrets.token_bytes(16)
    client_credentials = Credentials("foo", "bar")
    device_auth_hash = transport_class.generate_auth_hash(device_credentials)

    async def _return_handshake1_response(url, params=None, data=None, *_, **__):
        nonlocal client_seed

        client_seed = data
        seed_auth_hash = _sha256(
            seed_auth_hash_calc(client_seed, server_seed, device_auth_hash)
        )
        return _mock_response(200, server_seed + seed_auth_hash)

    mocker.patch.object(
        aiohttp.ClientSession, "post", side_effect=_return_handshake1_response
    )

    config = DeviceConfig("127.0.0.1", credentials=client_credentials)
    transport = transport_class(config=config)
    protocol = IotProtocol(transport=transport)

    with expectation:
        (
            local_seed,
            device_remote_seed,
            auth_hash,
        ) = await transport.perform_handshake1()

        assert local_seed == client_seed
        assert device_remote_seed == server_seed
        assert device_auth_hash == auth_hash
    await protocol.close()


@pytest.mark.parametrize(
    ("transport_class", "seed_auth_hash_calc1", "seed_auth_hash_calc2"),
    [
        pytest.param(
            KlapTransport, lambda c, s, a: c + a, lambda c, s, a: s + a, id="KLAP"
        ),
        pytest.param(
            KlapTransportV2,
            lambda c, s, a: c + s + a,
            lambda c, s, a: s + c + a,
            id="KLAPV2",
        ),
    ],
)
async def test_handshake(
    mocker: MockerFixture,
    transport_class: type[KlapTransport],
    seed_auth_hash_calc1: Callable[..., bytes],
    seed_auth_hash_calc2: Callable[..., bytes],
) -> None:
    client_seed = None
    server_seed = secrets.token_bytes(16)
    client_credentials = Credentials("foo", "bar")
    device_auth_hash = transport_class.generate_auth_hash(client_credentials)

    async def _return_handshake_response(url: URL, params=None, data=None, *_, **__):
        nonlocal client_seed

        if url == URL("http://127.0.0.1:80/app/handshake1"):
            client_seed = data
            seed_auth_hash = _sha256(
                seed_auth_hash_calc1(client_seed, server_seed, device_auth_hash)
            )

            return _mock_response(200, server_seed + seed_auth_hash)
        elif url == URL("http://127.0.0.1:80/app/handshake2"):
            seed_auth_hash = _sha256(
                seed_auth_hash_calc2(client_seed, server_seed, device_auth_hash)
            )
            assert data == seed_auth_hash
            return _mock_response(response_status, b"")

    mocker.patch.object(
        aiohttp.ClientSession, "post", side_effect=_return_handshake_response
    )

    config = DeviceConfig("127.0.0.1", credentials=client_credentials)
    transport = transport_class(config=config)
    protocol = IotProtocol(transport=transport)

    response_status = 200
    await transport.perform_handshake()
    assert transport._handshake_done is True

    response_status = 403
    with pytest.raises(KasaException):
        await transport.perform_handshake()
    assert transport._handshake_done is False
    await protocol.close()


async def test_query(mocker: MockerFixture) -> None:
    client_seed = None
    last_seq = None
    seq = None
    server_seed = secrets.token_bytes(16)
    client_credentials = Credentials("foo", "bar")
    device_auth_hash = KlapTransport.generate_auth_hash(client_credentials)

    async def _return_response(url: URL, params=None, data=None, *_, **__):
        nonlocal client_seed, seq

        if url == URL("http://127.0.0.1:80/app/handshake1"):
            client_seed = data
            client_seed_auth_hash = _sha256(data + device_auth_hash)

            return _mock_response(200, server_seed + client_seed_auth_hash)
        elif url == URL("http://127.0.0.1:80/app/handshake2"):
            return _mock_response(200, b"")
        elif url == URL("http://127.0.0.1:80/app/request"):
            assert transport._encryption_session is not None
            encryption_session = KlapEncryptionSession(
                transport._encryption_session.local_seed,
                transport._encryption_session.remote_seed,
                transport._encryption_session.user_hash,
            )
            seq = params.get("seq")
            encryption_session._seq = seq - 1
            encrypted, seq = encryption_session.encrypt('{"great": "success"}')
            seq = seq
            return _mock_response(200, encrypted)

    mocker.patch.object(aiohttp.ClientSession, "post", side_effect=_return_response)

    config = DeviceConfig("127.0.0.1", credentials=client_credentials)
    transport = KlapTransport(config=config)
    protocol = IotProtocol(transport=transport)

    for _ in range(10):
        resp = await protocol.query({})
        assert resp == {"great": "success"}
        # Check the protocol is incrementing the sequence number
        assert last_seq is None or last_seq + 1 == seq
        last_seq = seq


@pytest.mark.parametrize(
    ("response_status", "credentials_match", "expectation"),
    [
        pytest.param(
            (403, 403, 403),
            True,
            pytest.raises(KasaException),
            id="handshake1-403-status",
        ),
        pytest.param(
            (200, 403, 403),
            True,
            pytest.raises(KasaException),
            id="handshake2-403-status",
        ),
        pytest.param(
            (200, 200, 403),
            True,
            pytest.raises(_RetryableError),
            id="request-403-status",
        ),
        pytest.param(
            (200, 200, 400),
            True,
            pytest.raises(KasaException),
            id="request-400-status",
        ),
        pytest.param(
            (200, 200, 200),
            False,
            pytest.raises(AuthenticationError),
            id="handshake1-wrong-auth",
        ),
        pytest.param(
            (200, 200, 200),
            secrets.token_bytes(16),
            pytest.raises(KasaException),
            id="handshake1-bad-auth-length",
        ),
    ],
)
async def test_authentication_failures(
    mocker: MockerFixture,
    response_status: tuple[int, ...],
    credentials_match: bool | bytes,
    expectation: AbstractContextManager,
) -> None:
    client_seed = None

    server_seed = secrets.token_bytes(16)
    client_credentials = Credentials("foo", "bar")
    device_credentials = (
        client_credentials if credentials_match else Credentials("bar", "foo")
    )
    device_auth_hash = KlapTransport.generate_auth_hash(device_credentials)

    async def _return_response(url: URL, params=None, data=None, *_, **__):
        nonlocal client_seed

        if url == URL("http://127.0.0.1:80/app/handshake1"):
            client_seed = data
            client_seed_auth_hash = _sha256(data + device_auth_hash)
            if credentials_match is not False and credentials_match is not True:
                client_seed_auth_hash += credentials_match
            return _mock_response(
                response_status[0], server_seed + client_seed_auth_hash
            )
        elif url == URL("http://127.0.0.1:80/app/handshake2"):
            client_seed = data
            client_seed_auth_hash = _sha256(data + device_auth_hash)
            return _mock_response(
                response_status[1], server_seed + client_seed_auth_hash
            )
        elif url == URL("http://127.0.0.1:80/app/request"):
            return _mock_response(response_status[2], b"")

    mocker.patch.object(aiohttp.ClientSession, "post", side_effect=_return_response)

    config = DeviceConfig("127.0.0.1", credentials=client_credentials)
    protocol = IotProtocol(transport=KlapTransport(config=config))

    with expectation:
        await protocol.query({})


async def test_port_override() -> None:
    """Test that port override sets the app_url."""
    host = "127.0.0.1"
    config = DeviceConfig(
        host, credentials=Credentials("foo", "bar"), port_override=12345
    )
    transport = KlapTransport(config=config)

    assert str(transport._app_url) == "http://127.0.0.1:12345/app"


@device_iot
async def test_max_device_response_size_klap(
    dev: IotDevice, mocker: MockerFixture
) -> None:
    """Test that devices other than bulbs use a smaller budget over KLAP."""
    is_bulb = isinstance(dev, IotBulb)
    assert dev.max_device_response_size == (4096 if is_bulb else 16 * 1024)

    mocker.patch.object(dev.config, "connection_type", KLAP_CONNECTION)
    expected = 4096 if is_bulb else 2048
    assert dev.max_device_response_size == expected
    for child in dev.children:
        assert child.max_device_response_size == expected


def test_max_device_response_size_before_update() -> None:
    """Test that the budget follows the connection settings before an update."""
    assert IotPlug("127.0.0.1").max_device_response_size == 16 * 1024

    config = DeviceConfig("127.0.0.1", connection_type=KLAP_CONNECTION)
    assert IotPlug("127.0.0.1", config=config).max_device_response_size == 2048


@ks200m_iot
async def test_ks200m_update_split_klap(
    dev: IotWallSwitch, mocker: MockerFixture
) -> None:
    """Check that a KS200M update over KLAP is split into smaller requests."""
    spy = mocker.spy(dev.protocol, "query")
    await dev.update()
    assert spy.call_count == 1

    mocker.patch.object(dev.config, "connection_type", KLAP_CONNECTION)
    spy.reset_mock()
    await dev.update()

    assert [sorted(call.args[0]) for call in spy.call_args_list] == [
        ["schedule", "system"],
        ["schedule"],
        ["anti_theft", "cnCloud", "smartlife.iot.LAS", "smartlife.iot.PIR", "time"],
    ]