From f55428549f86648e1d25d44f5196c50075e8f54a Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sun, 19 Nov 2023 10:47:21 -0600
Subject: [PATCH] tweaks

---
 kasa/device_factory.py            | 25 +++++++++++++++++++++++++
 kasa/tests/test_device_factory.py | 29 ++++++++++++++++-------------
 kasa/tests/test_device_type.py    | 12 ------------
 3 files changed, 41 insertions(+), 25 deletions(-)

diff --git a/kasa/device_factory.py b/kasa/device_factory.py
index b2de926b..e0424b5e 100755
--- a/kasa/device_factory.py
+++ b/kasa/device_factory.py
@@ -1,5 +1,7 @@
 """Device creation by type."""
 
+import logging
+import time
 from typing import Any, Dict, Optional, Type
 
 from .credentials import Credentials
@@ -19,6 +21,8 @@ DEVICE_TYPE_TO_CLASS = {
     DeviceType.LightStrip: SmartLightStrip,
 }
 
+_LOGGER = logging.getLogger(__name__)
+
 
 async def connect(
     host: str,
@@ -48,11 +52,24 @@ async def connect(
     :rtype: SmartDevice
     :return: Object for querying/controlling found device.
     """
+    debug_enabled = _LOGGER.isEnabledFor(logging.DEBUG)
+
+    if debug_enabled:
+        start_time = time.perf_counter()
+
     if device_type and (klass := DEVICE_TYPE_TO_CLASS.get(device_type)):
         dev: SmartDevice = klass(
             host=host, port=port, credentials=credentials, timeout=timeout
         )
         await dev.update()
+        if debug_enabled:
+            end_time = time.perf_counter()
+            _LOGGER.debug(
+                "Device %s with known type (%s) took %.2f seconds to connect",
+                host,
+                device_type.value,
+                end_time - start_time,
+            )
         return dev
 
     unknown_dev = SmartDevice(
@@ -65,6 +82,14 @@ async def connect(
     # so we don't have to reconnect
     dev.protocol = unknown_dev.protocol
     await dev.update()
+    if debug_enabled:
+        end_time = time.perf_counter()
+        _LOGGER.debug(
+            "Device %s with unknown type (%s) took %.2f seconds to connect",
+            host,
+            dev.device_type.value,
+            end_time - start_time,
+        )
     return dev
 
 
diff --git a/kasa/tests/test_device_factory.py b/kasa/tests/test_device_factory.py
index ec03e5bc..3a08857a 100644
--- a/kasa/tests/test_device_factory.py
+++ b/kasa/tests/test_device_factory.py
@@ -1,32 +1,24 @@
 # type: ignore
-import re
-import socket
-import sys
+import logging
 from typing import Type
 
 import pytest  # type: ignore # https://github.com/pytest-dev/pytest/issues/3342
 
 from kasa import (
     DeviceType,
-    Discover,
     SmartBulb,
     SmartDevice,
     SmartDeviceException,
     SmartDimmer,
     SmartLightStrip,
     SmartPlug,
-    protocol,
 )
-from kasa.device_factory import DEVICE_TYPE_TO_CLASS, connect
-from kasa.discover import _DiscoverProtocol, json_dumps
-from kasa.exceptions import UnsupportedDeviceException
-
-from .conftest import bulb, dimmer, lightstrip, plug, strip
+from kasa.device_factory import connect
 
 
 @pytest.mark.parametrize("custom_port", [123, None])
 async def test_connect(discovery_data: dict, mocker, custom_port):
-    """Make sure that connect_single returns an initialized SmartDevice instance."""
+    """Make sure that connect returns an initialized SmartDevice instance."""
     host = "127.0.0.1"
     mocker.patch("kasa.TPLinkSmartHomeProtocol.query", return_value=discovery_data)
 
@@ -53,7 +45,7 @@ async def test_connect_passed_device_type(
     klass: Type[SmartDevice],
     custom_port,
 ):
-    """Make sure that connect_single with a 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)
 
@@ -63,9 +55,20 @@ async def test_connect_passed_device_type(
 
 
 async def test_connect_query_fails(discovery_data: dict, mocker):
-    """Make sure that connect_single fails when query fails."""
+    """Make sure that connect fails when query fails."""
     host = "127.0.0.1"
     mocker.patch("kasa.TPLinkSmartHomeProtocol.query", side_effect=SmartDeviceException)
 
     with pytest.raises(SmartDeviceException):
         await connect(host)
+
+
+async def test_connect_logs_connect_time(
+    discovery_data: dict, caplog: pytest.LogCaptureFixture, mocker
+):
+    """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
diff --git a/kasa/tests/test_device_type.py b/kasa/tests/test_device_type.py
index 8058d27e..da1707dc 100644
--- a/kasa/tests/test_device_type.py
+++ b/kasa/tests/test_device_type.py
@@ -1,16 +1,4 @@
-import inspect
-from datetime import datetime
-from unittest.mock import patch
-
-import pytest  # type: ignore # https://github.com/pytest-dev/pytest/issues/3342
-
-import kasa
-from kasa import Credentials, SmartDevice, SmartDeviceException
 from kasa.smartdevice import DeviceType
-from kasa.smartstrip import SmartStripPlug
-
-from .conftest import handle_turn_on, has_emeter, no_emeter, turn_on
-from .newfakes import PLUG_SCHEMA, TZ_SCHEMA, FakeTransportProtocol
 
 
 async def test_device_type_from_value():