Get the legacy_class from the config files for lock devices.

custom_components/tuya_local/helpers/device_config.py

@@ -115,7 +115,7 @@ class TuyaEntityConfig:
     @property
     def legacy_class(self):
         """Return the legacy device corresponding to this config."""
-        return self.__config.get("legacy_class", None)
+        return "custom_components.tuya_local" + self.__config.get("legacy_class", None)
 
     @property
     def entity(self):
@@ -204,3 +204,18 @@ def possible_matches(dps):
         parsed = TuyaDeviceConfig(cfg)
         if parsed.matches(dps):
             yield parsed
+
+
+def config_for_legacy_use(conf_type):
+    """
+    Return a config to use with config_type for legacy transition.
+    Note: as there are two variants for Kogan Socket, this is not guaranteed
+    to be the correct config for the device, so only use it for looking up
+    the legacy class during the transition period.
+    """
+    for cfg in available_configs():
+        parsed = TuyaDeviceConfig(cfg)
+        if parsed.legacy_type == conf_type:
+            return parsed
+
+    return None

custom_components/tuya_local/lock.py

@@ -2,6 +2,7 @@
 Setup for different kinds of Tuya climate devices
 """
 import logging
+from pydoc import locate
 
 from . import DOMAIN
 from .const import (
@@ -9,17 +10,8 @@ from .const import (
     CONF_DEVICE_ID,
     CONF_TYPE,
     CONF_TYPE_AUTO,
-    CONF_TYPE_DEHUMIDIFIER,
-    CONF_TYPE_GECO_HEATER,
-    CONF_TYPE_GPCV_HEATER,
-    CONF_TYPE_GPPH_HEATER,
-    CONF_TYPE_KOGAN_HEATER,
 )
-from .dehumidifier.lock import GoldairDehumidifierChildLock
-from .geco_heater.lock import GoldairGECOHeaterChildLock
-from .gpcv_heater.lock import GoldairGPCVHeaterChildLock
-from .heater.lock import GoldairHeaterChildLock
-from .kogan_heater.lock import KoganHeaterChildLock
+from .helpers.device_config import config_for_legacy_use
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -36,18 +28,20 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
         if discovery_info[CONF_TYPE] is None:
             raise ValueError(f"Unable to detect type for device {device.name}")
 
-    if discovery_info[CONF_TYPE] == CONF_TYPE_GPPH_HEATER:
-        data[CONF_CHILD_LOCK] = GoldairHeaterChildLock(device)
-    elif discovery_info[CONF_TYPE] == CONF_TYPE_DEHUMIDIFIER:
-        data[CONF_CHILD_LOCK] = GoldairDehumidifierChildLock(device)
-    elif discovery_info[CONF_TYPE] == CONF_TYPE_GECO_HEATER:
-        data[CONF_CHILD_LOCK] = GoldairGECOHeaterChildLock(device)
-    elif discovery_info[CONF_TYPE] == CONF_TYPE_GPCV_HEATER:
-        data[CONF_CHILD_LOCK] = GoldairGPCVHeaterChildLock(device)
-    elif discovery_info[CONF_TYPE] == CONF_TYPE_KOGAN_HEATER:
-        data[CONF_CHILD_LOCK] = KoganHeaterChildLock(device)
-    else:
-        raise ValueError("This device does not support child lock.")
+    cfg = config_for_legacy_use(discovery_info[CONF_TYPE])
+    ecfg = cfg.primary_entity
+    if ecfg.entity != "lock":
+        for ecfg in cfg.secondary_entities():
+            if ecfg.entity == "lock":
+                break
+        if ecfg.entity != "lock":
+            raise ValueError(f"{device.name} does not support use as a lock device.")
+
+    legacy_class = locate(ecfg.legacy_class)
+    if legacy_class is None:
+        raise TypeError(f"No legacy class {ecfg.legacy_class} exists.")
+
+    data[CONF_CHILD_LOCK] = legacy_class(device)
 
     if CONF_CHILD_LOCK in data:
         async_add_entities([data[CONF_CHILD_LOCK]])

requirements-dev.txt

@@ -3,6 +3,7 @@ homeassistant
 isort
 pytest-homeassistant-custom-component
 pytest
+pytest-asyncio
 pytest-cov
 pycryptodome~=3.9
 tinytuya~=1.2

tests/test_device_config.py

@@ -19,6 +19,7 @@ from custom_components.tuya_local.const import (
 
 from custom_components.tuya_local.helpers.device_config import (
     available_configs,
+    config_for_legacy_use,
     possible_matches,
     TuyaDeviceConfig,
 )
@@ -81,7 +82,10 @@ class TestDeviceConfig(unittest.TestCase):
                 self.assertFalse(matched)
                 matched = True
                 quality = cfg.match_quality(payload)
-                self.assertEqual(cfg.primary_entity.legacy_class, legacy_class)
+                self.assertEqual(
+                    cfg.primary_entity.legacy_class,
+                    "custom_components.tuya_local" + legacy_class,
+                )
             else:
                 false_matches.append(cfg)
 
@@ -98,6 +102,13 @@ class TestDeviceConfig(unittest.TestCase):
 
         self.assertGreater(quality, best_q)
 
+        # Ensure the same correct config is returned when looked up by type
+        cfg = config_for_legacy_use(legacy_type)
+        self.assertEqual(
+            cfg.primary_entity.legacy_class,
+            "custom_components.tuya_local" + legacy_class,
+        )
+
     def test_gpph_heater_detection(self):
         """Test that GPPH heater can be detected from its sample payload."""
         self._test_detect(

tests/test_lock.py

@@ -0,0 +1,34 @@
+"""Tests for the lock entity."""
+import pytest
+from pytest_homeassistant_custom_component.common import MockConfigEntry
+from unittest.mock import AsyncMock
+
+from custom_components.tuya_local.const import (
+    CONF_CHILD_LOCK,
+    CONF_DEVICE_ID,
+    CONF_TYPE,
+    CONF_TYPE_AUTO,
+    CONF_TYPE_GPPH_HEATER,
+    DOMAIN,
+)
+from custom_components.tuya_local.heater.lock import GoldairHeaterChildLock
+from custom_components.tuya_local.lock import async_setup_entry
+
+
+@pytest.mark.asyncio
+async def test_init_entry(hass):
+    """Test the initialisation."""
+    entry = MockConfigEntry(
+        domain=DOMAIN, data={CONF_TYPE: CONF_TYPE_AUTO, CONF_DEVICE_ID: "dummy"},
+    )
+    m_add_entities = AsyncMock()
+    m_device = AsyncMock()
+    m_device.async_inferred_type = AsyncMock(return_value=CONF_TYPE_GPPH_HEATER)
+
+    hass.data[DOMAIN] = {}
+    hass.data[DOMAIN]["dummy"] = {}
+    hass.data[DOMAIN]["dummy"]["device"] = m_device
+
+    await async_setup_entry(hass, entry, m_add_entities)
+    assert type(hass.data[DOMAIN]["dummy"][CONF_CHILD_LOCK]) == GoldairHeaterChildLock
+    m_add_entities.assert_called_once()