"""Support for VeSync switches.""" import logging from homeassistant.components.switch import SwitchEntity from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity_platform import AddEntitiesCallback from .common import VeSyncDevice from .const import DOMAIN, VS_DISCOVERY, VS_SWITCHES _LOGGER = logging.getLogger(__name__) DEV_TYPE_TO_HA = { "wifi-switch-1.3": "outlet", "ESW03-USA": "outlet", "ESW01-EU": "outlet", "ESW15-USA": "outlet", "ESWL01": "switch", "ESWL03": "switch", "ESO15-TB": "outlet", } async def async_setup_entry( hass: HomeAssistant, config_entry: ConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up switches.""" @callback def discover(devices): """Add new devices to platform.""" _setup_entities(devices, async_add_entities) config_entry.async_on_unload( async_dispatcher_connect(hass, VS_DISCOVERY.format(VS_SWITCHES), discover) ) _setup_entities(hass.data[DOMAIN][VS_SWITCHES], async_add_entities) @callback def _setup_entities(devices, async_add_entities): """Check if device is online and add entity.""" entities = [] for dev in devices: if DEV_TYPE_TO_HA.get(dev.device_type) == "outlet": entities.append(VeSyncSwitchHA(dev)) elif DEV_TYPE_TO_HA.get(dev.device_type) == "switch": entities.append(VeSyncLightSwitch(dev)) else: _LOGGER.warning( "%s - Unknown device type - %s", dev.device_name, dev.device_type ) continue async_add_entities(entities, update_before_add=True) class VeSyncBaseSwitch(VeSyncDevice, SwitchEntity): """Base class for VeSync switch Device Representations.""" def turn_on(self, **kwargs): """Turn the device on.""" self.device.turn_on() class VeSyncSwitchHA(VeSyncBaseSwitch, SwitchEntity): """Representation of a VeSync switch.""" def __init__(self, plug): """Initialize the VeSync switch device.""" super().__init__(plug) self.smartplug = plug @property def extra_state_attributes(self): """Return the state attributes of the device.""" if not hasattr(self.smartplug, "weekly_energy_total"): return {} return { "voltage": self.smartplug.voltage, "weekly_energy_total": self.smartplug.weekly_energy_total, "monthly_energy_total": self.smartplug.monthly_energy_total, "yearly_energy_total": self.smartplug.yearly_energy_total, } def update(self): """Update outlet details and energy usage.""" self.smartplug.update() self.smartplug.update_energy() class VeSyncLightSwitch(VeSyncBaseSwitch, SwitchEntity): """Handle representation of VeSync Light Switch.""" def __init__(self, switch): """Initialize Light Switch device class.""" super().__init__(switch) self.switch = switch