diff --git a/homeassistant/components/airgradient/__init__.py b/homeassistant/components/airgradient/__init__.py index 7ee8ac6a3c7..3b27d6cda5e 100644 --- a/homeassistant/components/airgradient/__init__.py +++ b/homeassistant/components/airgradient/__init__.py @@ -2,18 +2,14 @@ from __future__ import annotations -from dataclasses import dataclass - -from airgradient import AirGradientClient, get_model_name +from airgradient import AirGradientClient from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_HOST, Platform from homeassistant.core import HomeAssistant -from homeassistant.helpers import device_registry as dr from homeassistant.helpers.aiohttp_client import async_get_clientsession -from .const import DOMAIN -from .coordinator import AirGradientConfigCoordinator, AirGradientMeasurementCoordinator +from .coordinator import AirGradientCoordinator PLATFORMS: list[Platform] = [ Platform.BUTTON, @@ -25,15 +21,7 @@ PLATFORMS: list[Platform] = [ ] -@dataclass -class AirGradientData: - """AirGradient data class.""" - - measurement: AirGradientMeasurementCoordinator - config: AirGradientConfigCoordinator - - -type AirGradientConfigEntry = ConfigEntry[AirGradientData] +type AirGradientConfigEntry = ConfigEntry[AirGradientCoordinator] async def async_setup_entry(hass: HomeAssistant, entry: AirGradientConfigEntry) -> bool: @@ -43,27 +31,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: AirGradientConfigEntry) entry.data[CONF_HOST], session=async_get_clientsession(hass) ) - measurement_coordinator = AirGradientMeasurementCoordinator(hass, client) - config_coordinator = AirGradientConfigCoordinator(hass, client) + coordinator = AirGradientCoordinator(hass, client) - await measurement_coordinator.async_config_entry_first_refresh() - await config_coordinator.async_config_entry_first_refresh() + await coordinator.async_config_entry_first_refresh() - device_registry = dr.async_get(hass) - device_registry.async_get_or_create( - config_entry_id=entry.entry_id, - identifiers={(DOMAIN, measurement_coordinator.serial_number)}, - manufacturer="AirGradient", - model=get_model_name(measurement_coordinator.data.model), - model_id=measurement_coordinator.data.model, - serial_number=measurement_coordinator.data.serial_number, - sw_version=measurement_coordinator.data.firmware_version, - ) - - entry.runtime_data = AirGradientData( - measurement=measurement_coordinator, - config=config_coordinator, - ) + entry.runtime_data = coordinator await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) diff --git a/homeassistant/components/airgradient/button.py b/homeassistant/components/airgradient/button.py index b59188ebdd4..32a9b5adedf 100644 --- a/homeassistant/components/airgradient/button.py +++ b/homeassistant/components/airgradient/button.py @@ -15,8 +15,9 @@ from homeassistant.core import HomeAssistant, callback from homeassistant.helpers import entity_registry as er from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import DOMAIN, AirGradientConfigEntry -from .coordinator import AirGradientConfigCoordinator +from . import AirGradientConfigEntry +from .const import DOMAIN +from .coordinator import AirGradientCoordinator from .entity import AirGradientEntity @@ -47,8 +48,8 @@ async def async_setup_entry( async_add_entities: AddEntitiesCallback, ) -> None: """Set up AirGradient button entities based on a config entry.""" - model = entry.runtime_data.measurement.data.model - coordinator = entry.runtime_data.config + coordinator = entry.runtime_data + model = coordinator.data.measures.model added_entities = False @@ -57,7 +58,7 @@ async def async_setup_entry( nonlocal added_entities if ( - coordinator.data.configuration_control is ConfigurationControl.LOCAL + coordinator.data.config.configuration_control is ConfigurationControl.LOCAL and not added_entities ): entities = [AirGradientButton(coordinator, CO2_CALIBRATION)] @@ -67,7 +68,8 @@ async def async_setup_entry( async_add_entities(entities) added_entities = True elif ( - coordinator.data.configuration_control is not ConfigurationControl.LOCAL + coordinator.data.config.configuration_control + is not ConfigurationControl.LOCAL and added_entities ): entity_registry = er.async_get(hass) @@ -87,11 +89,10 @@ class AirGradientButton(AirGradientEntity, ButtonEntity): """Defines an AirGradient button.""" entity_description: AirGradientButtonEntityDescription - coordinator: AirGradientConfigCoordinator def __init__( self, - coordinator: AirGradientConfigCoordinator, + coordinator: AirGradientCoordinator, description: AirGradientButtonEntityDescription, ) -> None: """Initialize airgradient button.""" diff --git a/homeassistant/components/airgradient/coordinator.py b/homeassistant/components/airgradient/coordinator.py index c3def0b1f33..4e1c335019c 100644 --- a/homeassistant/components/airgradient/coordinator.py +++ b/homeassistant/components/airgradient/coordinator.py @@ -2,6 +2,7 @@ from __future__ import annotations +from dataclasses import dataclass from datetime import timedelta from typing import TYPE_CHECKING @@ -16,7 +17,15 @@ if TYPE_CHECKING: from . import AirGradientConfigEntry -class AirGradientCoordinator[_DataT](DataUpdateCoordinator[_DataT]): +@dataclass +class AirGradientData: + """Class for AirGradient data.""" + + measures: Measures + config: Config + + +class AirGradientCoordinator(DataUpdateCoordinator[AirGradientData]): """Class to manage fetching AirGradient data.""" config_entry: AirGradientConfigEntry @@ -33,25 +42,11 @@ class AirGradientCoordinator[_DataT](DataUpdateCoordinator[_DataT]): assert self.config_entry.unique_id self.serial_number = self.config_entry.unique_id - async def _async_update_data(self) -> _DataT: + async def _async_update_data(self) -> AirGradientData: try: - return await self._update_data() + measures = await self.client.get_current_measures() + config = await self.client.get_config() except AirGradientError as error: raise UpdateFailed(error) from error - - async def _update_data(self) -> _DataT: - raise NotImplementedError - - -class AirGradientMeasurementCoordinator(AirGradientCoordinator[Measures]): - """Class to manage fetching AirGradient data.""" - - async def _update_data(self) -> Measures: - return await self.client.get_current_measures() - - -class AirGradientConfigCoordinator(AirGradientCoordinator[Config]): - """Class to manage fetching AirGradient data.""" - - async def _update_data(self) -> Config: - return await self.client.get_config() + else: + return AirGradientData(measures, config) diff --git a/homeassistant/components/airgradient/entity.py b/homeassistant/components/airgradient/entity.py index 4de07904bba..588a799610b 100644 --- a/homeassistant/components/airgradient/entity.py +++ b/homeassistant/components/airgradient/entity.py @@ -1,5 +1,7 @@ """Base class for AirGradient entities.""" +from airgradient import get_model_name + from homeassistant.helpers.device_registry import DeviceInfo from homeassistant.helpers.update_coordinator import CoordinatorEntity @@ -15,6 +17,12 @@ class AirGradientEntity(CoordinatorEntity[AirGradientCoordinator]): def __init__(self, coordinator: AirGradientCoordinator) -> None: """Initialize airgradient entity.""" super().__init__(coordinator) + measures = coordinator.data.measures self._attr_device_info = DeviceInfo( identifiers={(DOMAIN, coordinator.serial_number)}, + manufacturer="AirGradient", + model=get_model_name(measures.model), + model_id=measures.model, + serial_number=coordinator.serial_number, + sw_version=measures.firmware_version, ) diff --git a/homeassistant/components/airgradient/number.py b/homeassistant/components/airgradient/number.py index 139357f3753..7fd282ddd8b 100644 --- a/homeassistant/components/airgradient/number.py +++ b/homeassistant/components/airgradient/number.py @@ -18,7 +18,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback from . import AirGradientConfigEntry from .const import DOMAIN -from .coordinator import AirGradientConfigCoordinator +from .coordinator import AirGradientCoordinator from .entity import AirGradientEntity @@ -62,8 +62,8 @@ async def async_setup_entry( ) -> None: """Set up AirGradient number entities based on a config entry.""" - model = entry.runtime_data.measurement.data.model - coordinator = entry.runtime_data.config + coordinator = entry.runtime_data + model = coordinator.data.measures.model added_entities = False @@ -72,7 +72,7 @@ async def async_setup_entry( nonlocal added_entities if ( - coordinator.data.configuration_control is ConfigurationControl.LOCAL + coordinator.data.config.configuration_control is ConfigurationControl.LOCAL and not added_entities ): entities = [] @@ -84,7 +84,8 @@ async def async_setup_entry( async_add_entities(entities) added_entities = True elif ( - coordinator.data.configuration_control is not ConfigurationControl.LOCAL + coordinator.data.config.configuration_control + is not ConfigurationControl.LOCAL and added_entities ): entity_registry = er.async_get(hass) @@ -104,11 +105,10 @@ class AirGradientNumber(AirGradientEntity, NumberEntity): """Defines an AirGradient number entity.""" entity_description: AirGradientNumberEntityDescription - coordinator: AirGradientConfigCoordinator def __init__( self, - coordinator: AirGradientConfigCoordinator, + coordinator: AirGradientCoordinator, description: AirGradientNumberEntityDescription, ) -> None: """Initialize AirGradient number.""" @@ -119,7 +119,7 @@ class AirGradientNumber(AirGradientEntity, NumberEntity): @property def native_value(self) -> int | None: """Return the state of the number.""" - return self.entity_description.value_fn(self.coordinator.data) + return self.entity_description.value_fn(self.coordinator.data.config) async def async_set_native_value(self, value: float) -> None: """Set the selected value.""" diff --git a/homeassistant/components/airgradient/select.py b/homeassistant/components/airgradient/select.py index 532f7167dff..af56802d842 100644 --- a/homeassistant/components/airgradient/select.py +++ b/homeassistant/components/airgradient/select.py @@ -18,7 +18,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback from . import AirGradientConfigEntry from .const import DOMAIN, PM_STANDARD, PM_STANDARD_REVERSE -from .coordinator import AirGradientConfigCoordinator +from .coordinator import AirGradientCoordinator from .entity import AirGradientEntity @@ -144,13 +144,11 @@ async def async_setup_entry( ) -> None: """Set up AirGradient select entities based on a config entry.""" - coordinator = entry.runtime_data.config - measurement_coordinator = entry.runtime_data.measurement + coordinator = entry.runtime_data + model = coordinator.data.measures.model async_add_entities([AirGradientSelect(coordinator, CONFIG_CONTROL_ENTITY)]) - model = measurement_coordinator.data.model - added_entities = False @callback @@ -158,7 +156,7 @@ async def async_setup_entry( nonlocal added_entities if ( - coordinator.data.configuration_control is ConfigurationControl.LOCAL + coordinator.data.config.configuration_control is ConfigurationControl.LOCAL and not added_entities ): entities: list[AirGradientSelect] = [ @@ -179,7 +177,8 @@ async def async_setup_entry( async_add_entities(entities) added_entities = True elif ( - coordinator.data.configuration_control is not ConfigurationControl.LOCAL + coordinator.data.config.configuration_control + is not ConfigurationControl.LOCAL and added_entities ): entity_registry = er.async_get(hass) @@ -201,11 +200,10 @@ class AirGradientSelect(AirGradientEntity, SelectEntity): """Defines an AirGradient select entity.""" entity_description: AirGradientSelectEntityDescription - coordinator: AirGradientConfigCoordinator def __init__( self, - coordinator: AirGradientConfigCoordinator, + coordinator: AirGradientCoordinator, description: AirGradientSelectEntityDescription, ) -> None: """Initialize AirGradient select.""" @@ -216,7 +214,7 @@ class AirGradientSelect(AirGradientEntity, SelectEntity): @property def current_option(self) -> str | None: """Return the state of the select.""" - return self.entity_description.value_fn(self.coordinator.data) + return self.entity_description.value_fn(self.coordinator.data.config) async def async_select_option(self, option: str) -> None: """Change the selected option.""" diff --git a/homeassistant/components/airgradient/sensor.py b/homeassistant/components/airgradient/sensor.py index f431c49ed2a..497d4cc0488 100644 --- a/homeassistant/components/airgradient/sensor.py +++ b/homeassistant/components/airgradient/sensor.py @@ -32,7 +32,7 @@ from homeassistant.helpers.typing import StateType from . import AirGradientConfigEntry from .const import PM_STANDARD, PM_STANDARD_REVERSE -from .coordinator import AirGradientConfigCoordinator, AirGradientMeasurementCoordinator +from .coordinator import AirGradientCoordinator from .entity import AirGradientEntity @@ -218,7 +218,7 @@ async def async_setup_entry( ) -> None: """Set up AirGradient sensor entities based on a config entry.""" - coordinator = entry.runtime_data.measurement + coordinator = entry.runtime_data listener: Callable[[], None] | None = None not_setup: set[AirGradientMeasurementSensorEntityDescription] = set( MEASUREMENT_SENSOR_TYPES @@ -232,7 +232,7 @@ async def async_setup_entry( not_setup = set() sensors = [] for description in sensor_descriptions: - if description.value_fn(coordinator.data) is None: + if description.value_fn(coordinator.data.measures) is None: not_setup.add(description) else: sensors.append(AirGradientMeasurementSensor(coordinator, description)) @@ -248,64 +248,65 @@ async def async_setup_entry( add_entities() entities = [ - AirGradientConfigSensor(entry.runtime_data.config, description) + AirGradientConfigSensor(coordinator, description) for description in CONFIG_SENSOR_TYPES ] - if "L" in coordinator.data.model: + if "L" in coordinator.data.measures.model: entities.extend( - AirGradientConfigSensor(entry.runtime_data.config, description) + AirGradientConfigSensor(coordinator, description) for description in CONFIG_LED_BAR_SENSOR_TYPES ) - if "I" in coordinator.data.model: + if "I" in coordinator.data.measures.model: entities.extend( - AirGradientConfigSensor(entry.runtime_data.config, description) + AirGradientConfigSensor(coordinator, description) for description in CONFIG_DISPLAY_SENSOR_TYPES ) async_add_entities(entities) -class AirGradientMeasurementSensor(AirGradientEntity, SensorEntity): +class AirGradientSensor(AirGradientEntity, SensorEntity): """Defines an AirGradient sensor.""" - entity_description: AirGradientMeasurementSensorEntityDescription - coordinator: AirGradientMeasurementCoordinator - def __init__( self, - coordinator: AirGradientMeasurementCoordinator, - description: AirGradientMeasurementSensorEntityDescription, + coordinator: AirGradientCoordinator, + description: SensorEntityDescription, ) -> None: """Initialize airgradient sensor.""" super().__init__(coordinator) self.entity_description = description self._attr_unique_id = f"{coordinator.serial_number}-{description.key}" + +class AirGradientMeasurementSensor(AirGradientSensor): + """Defines an AirGradient sensor.""" + + entity_description: AirGradientMeasurementSensorEntityDescription + @property def native_value(self) -> StateType: """Return the state of the sensor.""" - return self.entity_description.value_fn(self.coordinator.data) + return self.entity_description.value_fn(self.coordinator.data.measures) -class AirGradientConfigSensor(AirGradientEntity, SensorEntity): +class AirGradientConfigSensor(AirGradientSensor): """Defines an AirGradient sensor.""" entity_description: AirGradientConfigSensorEntityDescription - coordinator: AirGradientConfigCoordinator def __init__( self, - coordinator: AirGradientConfigCoordinator, + coordinator: AirGradientCoordinator, description: AirGradientConfigSensorEntityDescription, ) -> None: """Initialize airgradient sensor.""" - super().__init__(coordinator) - self.entity_description = description - self._attr_unique_id = f"{coordinator.serial_number}-{description.key}" + super().__init__(coordinator, description) self._attr_entity_registry_enabled_default = ( - coordinator.data.configuration_control is not ConfigurationControl.LOCAL + coordinator.data.config.configuration_control + is not ConfigurationControl.LOCAL ) @property def native_value(self) -> StateType: """Return the state of the sensor.""" - return self.entity_description.value_fn(self.coordinator.data) + return self.entity_description.value_fn(self.coordinator.data.config) diff --git a/homeassistant/components/airgradient/switch.py b/homeassistant/components/airgradient/switch.py index 60c3f83ae5e..329f704e755 100644 --- a/homeassistant/components/airgradient/switch.py +++ b/homeassistant/components/airgradient/switch.py @@ -19,7 +19,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback from . import AirGradientConfigEntry from .const import DOMAIN -from .coordinator import AirGradientConfigCoordinator +from .coordinator import AirGradientCoordinator from .entity import AirGradientEntity @@ -46,7 +46,7 @@ async def async_setup_entry( async_add_entities: AddEntitiesCallback, ) -> None: """Set up AirGradient switch entities based on a config entry.""" - coordinator = entry.runtime_data.config + coordinator = entry.runtime_data added_entities = False @@ -55,7 +55,7 @@ async def async_setup_entry( nonlocal added_entities if ( - coordinator.data.configuration_control is ConfigurationControl.LOCAL + coordinator.data.config.configuration_control is ConfigurationControl.LOCAL and not added_entities ): async_add_entities( @@ -63,7 +63,8 @@ async def async_setup_entry( ) added_entities = True elif ( - coordinator.data.configuration_control is not ConfigurationControl.LOCAL + coordinator.data.config.configuration_control + is not ConfigurationControl.LOCAL and added_entities ): entity_registry = er.async_get(hass) @@ -82,11 +83,10 @@ class AirGradientSwitch(AirGradientEntity, SwitchEntity): """Defines an AirGradient switch entity.""" entity_description: AirGradientSwitchEntityDescription - coordinator: AirGradientConfigCoordinator def __init__( self, - coordinator: AirGradientConfigCoordinator, + coordinator: AirGradientCoordinator, description: AirGradientSwitchEntityDescription, ) -> None: """Initialize AirGradient switch.""" @@ -97,7 +97,7 @@ class AirGradientSwitch(AirGradientEntity, SwitchEntity): @property def is_on(self) -> bool: """Return the state of the switch.""" - return self.entity_description.value_fn(self.coordinator.data) + return self.entity_description.value_fn(self.coordinator.data.config) async def async_turn_on(self, **kwargs: Any) -> None: """Turn the switch on.""" diff --git a/homeassistant/components/airgradient/update.py b/homeassistant/components/airgradient/update.py index 95e64930ea6..eb6708afb67 100644 --- a/homeassistant/components/airgradient/update.py +++ b/homeassistant/components/airgradient/update.py @@ -7,7 +7,7 @@ from homeassistant.components.update import UpdateDeviceClass, UpdateEntity from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import AirGradientConfigEntry, AirGradientMeasurementCoordinator +from . import AirGradientConfigEntry, AirGradientCoordinator from .entity import AirGradientEntity SCAN_INTERVAL = timedelta(hours=1) @@ -20,18 +20,17 @@ async def async_setup_entry( ) -> None: """Set up Airgradient update platform.""" - data = config_entry.runtime_data + coordinator = config_entry.runtime_data - async_add_entities([AirGradientUpdate(data.measurement)], True) + async_add_entities([AirGradientUpdate(coordinator)], True) class AirGradientUpdate(AirGradientEntity, UpdateEntity): """Representation of Airgradient Update.""" _attr_device_class = UpdateDeviceClass.FIRMWARE - coordinator: AirGradientMeasurementCoordinator - def __init__(self, coordinator: AirGradientMeasurementCoordinator) -> None: + def __init__(self, coordinator: AirGradientCoordinator) -> None: """Initialize the entity.""" super().__init__(coordinator) self._attr_unique_id = f"{coordinator.serial_number}-update" @@ -44,7 +43,7 @@ class AirGradientUpdate(AirGradientEntity, UpdateEntity): @property def installed_version(self) -> str: """Return the installed version of the entity.""" - return self.coordinator.data.firmware_version + return self.coordinator.data.measures.firmware_version async def async_update(self) -> None: """Update the entity.""" diff --git a/tests/components/airgradient/snapshots/test_init.ambr b/tests/components/airgradient/snapshots/test_init.ambr index e47c5b38bbc..72cb12535f1 100644 --- a/tests/components/airgradient/snapshots/test_init.ambr +++ b/tests/components/airgradient/snapshots/test_init.ambr @@ -57,7 +57,7 @@ 'name': 'Airgradient', 'name_by_user': None, 'primary_config_entry': , - 'serial_number': '84fce60bec38', + 'serial_number': '84fce612f5b8', 'suggested_area': None, 'sw_version': '3.1.1', 'via_device_id': None, diff --git a/tests/components/airgradient/test_button.py b/tests/components/airgradient/test_button.py index 7901c3a067b..83de2c2f048 100644 --- a/tests/components/airgradient/test_button.py +++ b/tests/components/airgradient/test_button.py @@ -7,7 +7,7 @@ from airgradient import Config from freezegun.api import FrozenDateTimeFactory from syrupy import SnapshotAssertion -from homeassistant.components.airgradient import DOMAIN +from homeassistant.components.airgradient.const import DOMAIN from homeassistant.components.button import DOMAIN as BUTTON_DOMAIN, SERVICE_PRESS from homeassistant.const import ATTR_ENTITY_ID, Platform from homeassistant.core import HomeAssistant diff --git a/tests/components/airgradient/test_config_flow.py b/tests/components/airgradient/test_config_flow.py index 8730b18676f..73dbd17a213 100644 --- a/tests/components/airgradient/test_config_flow.py +++ b/tests/components/airgradient/test_config_flow.py @@ -9,7 +9,7 @@ from airgradient import ( ConfigurationControl, ) -from homeassistant.components.airgradient import DOMAIN +from homeassistant.components.airgradient.const import DOMAIN from homeassistant.components.zeroconf import ZeroconfServiceInfo from homeassistant.config_entries import SOURCE_USER, SOURCE_ZEROCONF from homeassistant.const import CONF_HOST diff --git a/tests/components/airgradient/test_init.py b/tests/components/airgradient/test_init.py index 408e6f5f3ba..a566254d106 100644 --- a/tests/components/airgradient/test_init.py +++ b/tests/components/airgradient/test_init.py @@ -4,7 +4,7 @@ from unittest.mock import AsyncMock from syrupy import SnapshotAssertion -from homeassistant.components.airgradient import DOMAIN +from homeassistant.components.airgradient.const import DOMAIN from homeassistant.core import HomeAssistant from homeassistant.helpers import device_registry as dr diff --git a/tests/components/airgradient/test_number.py b/tests/components/airgradient/test_number.py index 0803c0d437f..7aabda8f81c 100644 --- a/tests/components/airgradient/test_number.py +++ b/tests/components/airgradient/test_number.py @@ -7,7 +7,7 @@ from airgradient import Config from freezegun.api import FrozenDateTimeFactory from syrupy import SnapshotAssertion -from homeassistant.components.airgradient import DOMAIN +from homeassistant.components.airgradient.const import DOMAIN from homeassistant.components.number import ( ATTR_VALUE, DOMAIN as NUMBER_DOMAIN, diff --git a/tests/components/airgradient/test_select.py b/tests/components/airgradient/test_select.py index 61679a15c07..de4a7beaaa7 100644 --- a/tests/components/airgradient/test_select.py +++ b/tests/components/airgradient/test_select.py @@ -8,7 +8,7 @@ from freezegun.api import FrozenDateTimeFactory import pytest from syrupy import SnapshotAssertion -from homeassistant.components.airgradient import DOMAIN +from homeassistant.components.airgradient.const import DOMAIN from homeassistant.components.select import ( DOMAIN as SELECT_DOMAIN, SERVICE_SELECT_OPTION, diff --git a/tests/components/airgradient/test_sensor.py b/tests/components/airgradient/test_sensor.py index c2e53ef4de2..e3fed70839a 100644 --- a/tests/components/airgradient/test_sensor.py +++ b/tests/components/airgradient/test_sensor.py @@ -8,7 +8,7 @@ from freezegun.api import FrozenDateTimeFactory import pytest from syrupy import SnapshotAssertion -from homeassistant.components.airgradient import DOMAIN +from homeassistant.components.airgradient.const import DOMAIN from homeassistant.const import STATE_UNAVAILABLE, Platform from homeassistant.core import HomeAssistant from homeassistant.helpers import entity_registry as er diff --git a/tests/components/airgradient/test_switch.py b/tests/components/airgradient/test_switch.py index 20a1cb7470b..a0cbdd17d75 100644 --- a/tests/components/airgradient/test_switch.py +++ b/tests/components/airgradient/test_switch.py @@ -7,7 +7,7 @@ from airgradient import Config from freezegun.api import FrozenDateTimeFactory from syrupy import SnapshotAssertion -from homeassistant.components.airgradient import DOMAIN +from homeassistant.components.airgradient.const import DOMAIN from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN from homeassistant.const import ( ATTR_ENTITY_ID,