diff --git a/custom_components/localtuya/__init__.py b/custom_components/localtuya/__init__.py
index ff07ef8..022bffb 100644
--- a/custom_components/localtuya/__init__.py
+++ b/custom_components/localtuya/__init__.py
@@ -36,7 +36,7 @@ BASE_PLATFORM_SCHEMA = {
 }
 
 
-def prepare_setup_entities(config_entry, platform):
+def prepare_setup_entities(config_entry, platform, device_class):
     """Prepare ro setup entities for a platform."""
     entities_to_setup = [
         entity
@@ -46,7 +46,7 @@ def prepare_setup_entities(config_entry, platform):
     if not entities_to_setup:
         return None, None
 
-    device = pytuya.BulbDevice(
+    device = device_class(
         config_entry.data[CONF_DEVICE_ID],
         config_entry.data[CONF_HOST],
         config_entry.data[CONF_LOCAL_KEY],
diff --git a/custom_components/localtuya/light.py b/custom_components/localtuya/light.py
index 37c57b0..7f68c51 100644
--- a/custom_components/localtuya/light.py
+++ b/custom_components/localtuya/light.py
@@ -34,6 +34,7 @@ from homeassistant.components.light import (
 from homeassistant.util import color as colorutil
 
 from . import BASE_PLATFORM_SCHEMA, import_from_yaml, prepare_setup_entities
+from .pytuya import BulbDevice
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -51,7 +52,9 @@ def flow_schema(dps):
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Setup a Tuya switch based on a config entry."""
-    device, entities_to_setup = prepare_setup_entities(config_entry, "light")
+    device, entities_to_setup = prepare_setup_entities(
+        config_entry, "light", BulbDevice
+    )
     if not entities_to_setup:
         return
 
diff --git a/custom_components/localtuya/switch.py b/custom_components/localtuya/switch.py
index 8b12b6c..9dd64da 100644
--- a/custom_components/localtuya/switch.py
+++ b/custom_components/localtuya/switch.py
@@ -48,6 +48,7 @@ from .const import (
     CONF_CURRENT_CONSUMPTION,
     CONF_VOLTAGE,
 )
+from .pytuya import OutletDevice
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -86,7 +87,9 @@ def flow_schema(dps):
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Setup a Tuya switch based on a config entry."""
-    device, entities_to_setup = prepare_setup_entities(config_entry, "switch")
+    device, entities_to_setup = prepare_setup_entities(
+        config_entry, "switch", OutletDevice
+    )
     if not entities_to_setup:
         return