diff --git a/homeassistant/components/octoprint/__init__.py b/homeassistant/components/octoprint/__init__.py index a582832d477..a6c45ef26a3 100644 --- a/homeassistant/components/octoprint/__init__.py +++ b/homeassistant/components/octoprint/__init__.py @@ -9,7 +9,7 @@ import aiohttp from pyoctoprintapi import OctoprintClient import voluptuous as vol -from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry +from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntryState from homeassistant.const import ( CONF_API_KEY, CONF_BINARY_SENSORS, @@ -34,7 +34,7 @@ from homeassistant.util import slugify as util_slugify from homeassistant.util.ssl import get_default_context, get_default_no_verify_context from .const import CONF_BAUDRATE, DOMAIN, SERVICE_CONNECT -from .coordinator import OctoprintDataUpdateCoordinator +from .coordinator import OctoprintConfigEntry, OctoprintDataUpdateCoordinator _LOGGER = logging.getLogger(__name__) @@ -168,12 +168,8 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: return True -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, entry: OctoprintConfigEntry) -> bool: """Set up OctoPrint from a config entry.""" - - if DOMAIN not in hass.data: - hass.data[DOMAIN] = {} - if CONF_VERIFY_SSL not in entry.data: data = {**entry.data, CONF_VERIFY_SSL: True} hass.config_entries.async_update_entry(entry, data=data) @@ -210,10 +206,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: await coordinator.async_config_entry_first_refresh() - hass.data[DOMAIN][entry.entry_id] = { - "coordinator": coordinator, - "client": client, - } + entry.runtime_data = coordinator await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) @@ -237,14 +230,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: return True -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_unload_entry(hass: HomeAssistant, entry: OctoprintConfigEntry) -> bool: """Unload a config entry.""" - unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS) - - if unload_ok: - hass.data[DOMAIN].pop(entry.entry_id) - - return unload_ok + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) def async_get_client_for_service_call( @@ -256,8 +244,9 @@ def async_get_client_for_service_call( if device_entry := device_registry.async_get(device_id): for entry_id in device_entry.config_entries: - if data := hass.data[DOMAIN].get(entry_id): - return cast(OctoprintClient, data["client"]) + if entry := hass.config_entries.async_get_entry(entry_id): + if entry.domain == DOMAIN and entry.state == ConfigEntryState.LOADED: + return cast(OctoprintConfigEntry, entry).runtime_data.octoprint raise ServiceValidationError( translation_domain=DOMAIN, diff --git a/homeassistant/components/octoprint/binary_sensor.py b/homeassistant/components/octoprint/binary_sensor.py index 4d12ef15a4e..deb3059458f 100644 --- a/homeassistant/components/octoprint/binary_sensor.py +++ b/homeassistant/components/octoprint/binary_sensor.py @@ -7,24 +7,20 @@ from abc import abstractmethod from pyoctoprintapi import OctoprintPrinterInfo from homeassistant.components.binary_sensor import BinarySensorEntity -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback from homeassistant.helpers.update_coordinator import CoordinatorEntity -from . import OctoprintDataUpdateCoordinator -from .const import DOMAIN +from .coordinator import OctoprintConfigEntry, OctoprintDataUpdateCoordinator async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: OctoprintConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up the available OctoPrint binary sensors.""" - coordinator: OctoprintDataUpdateCoordinator = hass.data[DOMAIN][ - config_entry.entry_id - ]["coordinator"] + coordinator = config_entry.runtime_data device_id = config_entry.unique_id assert device_id is not None diff --git a/homeassistant/components/octoprint/button.py b/homeassistant/components/octoprint/button.py index 3a128fcd7aa..fe167702745 100644 --- a/homeassistant/components/octoprint/button.py +++ b/homeassistant/components/octoprint/button.py @@ -3,26 +3,22 @@ from pyoctoprintapi import OctoprintClient, OctoprintPrinterInfo from homeassistant.components.button import ButtonDeviceClass, ButtonEntity -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback from homeassistant.helpers.update_coordinator import CoordinatorEntity -from . import OctoprintDataUpdateCoordinator -from .const import DOMAIN +from .coordinator import OctoprintConfigEntry, OctoprintDataUpdateCoordinator async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: OctoprintConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up Octoprint control buttons.""" - coordinator: OctoprintDataUpdateCoordinator = hass.data[DOMAIN][ - config_entry.entry_id - ]["coordinator"] - client: OctoprintClient = hass.data[DOMAIN][config_entry.entry_id]["client"] + coordinator = config_entry.runtime_data + client = coordinator.octoprint device_id = config_entry.unique_id assert device_id is not None diff --git a/homeassistant/components/octoprint/camera.py b/homeassistant/components/octoprint/camera.py index 37347539d5b..118f892ed5b 100644 --- a/homeassistant/components/octoprint/camera.py +++ b/homeassistant/components/octoprint/camera.py @@ -2,29 +2,25 @@ from __future__ import annotations -from pyoctoprintapi import OctoprintClient, WebcamSettings +from pyoctoprintapi import WebcamSettings from homeassistant.components.mjpeg import MjpegCamera -from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_VERIFY_SSL from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback from homeassistant.helpers.update_coordinator import CoordinatorEntity -from . import OctoprintDataUpdateCoordinator -from .const import DOMAIN +from .coordinator import OctoprintConfigEntry, OctoprintDataUpdateCoordinator async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: OctoprintConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up the available OctoPrint camera.""" - coordinator: OctoprintDataUpdateCoordinator = hass.data[DOMAIN][ - config_entry.entry_id - ]["coordinator"] - client: OctoprintClient = hass.data[DOMAIN][config_entry.entry_id]["client"] + coordinator = config_entry.runtime_data + client = coordinator.octoprint device_id = config_entry.unique_id assert device_id is not None diff --git a/homeassistant/components/octoprint/coordinator.py b/homeassistant/components/octoprint/coordinator.py index bb006329ff1..f37fbc82f54 100644 --- a/homeassistant/components/octoprint/coordinator.py +++ b/homeassistant/components/octoprint/coordinator.py @@ -20,19 +20,21 @@ from homeassistant.util import dt as dt_util from .const import DOMAIN +type OctoprintConfigEntry = ConfigEntry[OctoprintDataUpdateCoordinator] + _LOGGER = logging.getLogger(__name__) class OctoprintDataUpdateCoordinator(DataUpdateCoordinator): """Class to manage fetching Octoprint data.""" - config_entry: ConfigEntry + config_entry: OctoprintConfigEntry def __init__( self, hass: HomeAssistant, octoprint: OctoprintClient, - config_entry: ConfigEntry, + config_entry: OctoprintConfigEntry, interval: int, ) -> None: """Initialize.""" @@ -43,7 +45,7 @@ class OctoprintDataUpdateCoordinator(DataUpdateCoordinator): name=f"octoprint-{config_entry.entry_id}", update_interval=timedelta(seconds=interval), ) - self._octoprint = octoprint + self.octoprint = octoprint self._printer_offline = False self.data = {"printer": None, "job": None, "last_read_time": None} @@ -51,7 +53,7 @@ class OctoprintDataUpdateCoordinator(DataUpdateCoordinator): """Update data via API.""" printer = None try: - job = await self._octoprint.get_job_info() + job = await self.octoprint.get_job_info() except UnauthorizedException as err: raise ConfigEntryAuthFailed from err except ApiError as err: @@ -61,7 +63,7 @@ class OctoprintDataUpdateCoordinator(DataUpdateCoordinator): # printer will return a 409, so continue using the last # reading if there is one try: - printer = await self._octoprint.get_printer_info() + printer = await self.octoprint.get_printer_info() except PrinterOffline: if not self._printer_offline: _LOGGER.debug("Unable to retrieve printer information: Printer offline") diff --git a/homeassistant/components/octoprint/number.py b/homeassistant/components/octoprint/number.py index 93fa32a9e33..abe27006dfd 100644 --- a/homeassistant/components/octoprint/number.py +++ b/homeassistant/components/octoprint/number.py @@ -7,15 +7,14 @@ import logging from pyoctoprintapi import OctoprintClient from homeassistant.components.number import NumberDeviceClass, NumberEntity, NumberMode -from homeassistant.config_entries import ConfigEntry from homeassistant.const import UnitOfTemperature from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback from homeassistant.helpers.update_coordinator import CoordinatorEntity -from . import OctoprintDataUpdateCoordinator from .const import DOMAIN +from .coordinator import OctoprintConfigEntry, OctoprintDataUpdateCoordinator _LOGGER = logging.getLogger(__name__) @@ -37,14 +36,12 @@ def is_first_extruder(tool_name: str) -> bool: async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: OctoprintConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up the OctoPrint number entities.""" - coordinator: OctoprintDataUpdateCoordinator = hass.data[DOMAIN][ - config_entry.entry_id - ]["coordinator"] - client: OctoprintClient = hass.data[DOMAIN][config_entry.entry_id]["client"] + coordinator = config_entry.runtime_data + client = coordinator.octoprint device_id = config_entry.unique_id assert device_id is not None diff --git a/homeassistant/components/octoprint/sensor.py b/homeassistant/components/octoprint/sensor.py index 26ef8721d51..485126b4828 100644 --- a/homeassistant/components/octoprint/sensor.py +++ b/homeassistant/components/octoprint/sensor.py @@ -12,14 +12,12 @@ from homeassistant.components.sensor import ( SensorEntity, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import PERCENTAGE, UnitOfInformation, UnitOfTemperature from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback from homeassistant.helpers.update_coordinator import CoordinatorEntity -from . import OctoprintDataUpdateCoordinator -from .const import DOMAIN +from .coordinator import OctoprintConfigEntry, OctoprintDataUpdateCoordinator _LOGGER = logging.getLogger(__name__) @@ -35,13 +33,11 @@ def _is_printer_printing(printer: OctoprintPrinterInfo) -> bool: async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: OctoprintConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up the available OctoPrint sensors.""" - coordinator: OctoprintDataUpdateCoordinator = hass.data[DOMAIN][ - config_entry.entry_id - ]["coordinator"] + coordinator = config_entry.runtime_data device_id = config_entry.unique_id assert device_id is not None diff --git a/tests/components/octoprint/test_button.py b/tests/components/octoprint/test_button.py index f3fff39dab1..4ebb25362da 100644 --- a/tests/components/octoprint/test_button.py +++ b/tests/components/octoprint/test_button.py @@ -8,10 +8,11 @@ import pytest from homeassistant.components.button import DOMAIN as BUTTON_DOMAIN, SERVICE_PRESS from homeassistant.components.octoprint import OctoprintDataUpdateCoordinator from homeassistant.components.octoprint.button import InvalidPrinterState -from homeassistant.components.octoprint.const import DOMAIN from homeassistant.const import ATTR_ENTITY_ID, Platform from homeassistant.core import HomeAssistant +from tests.common import MockConfigEntry + @pytest.fixture def platform() -> Platform: @@ -19,12 +20,11 @@ def platform() -> Platform: return Platform.BUTTON -@pytest.mark.usefixtures("init_integration") -async def test_pause_job(hass: HomeAssistant) -> None: +async def test_pause_job( + hass: HomeAssistant, init_integration: MockConfigEntry +) -> None: """Test the pause job button.""" - coordinator: OctoprintDataUpdateCoordinator = hass.data[DOMAIN]["uuid"][ - "coordinator" - ] + coordinator: OctoprintDataUpdateCoordinator = init_integration.runtime_data # Test pausing the printer when it is printing with patch("pyoctoprintapi.OctoprintClient.pause_job") as pause_command: @@ -77,12 +77,11 @@ async def test_pause_job(hass: HomeAssistant) -> None: ) -@pytest.mark.usefixtures("init_integration") -async def test_resume_job(hass: HomeAssistant) -> None: +async def test_resume_job( + hass: HomeAssistant, init_integration: MockConfigEntry +) -> None: """Test the resume job button.""" - coordinator: OctoprintDataUpdateCoordinator = hass.data[DOMAIN]["uuid"][ - "coordinator" - ] + coordinator: OctoprintDataUpdateCoordinator = init_integration.runtime_data # Test resuming the printer when it is paused with patch("pyoctoprintapi.OctoprintClient.resume_job") as resume_command: @@ -135,12 +134,9 @@ async def test_resume_job(hass: HomeAssistant) -> None: ) -@pytest.mark.usefixtures("init_integration") -async def test_stop_job(hass: HomeAssistant) -> None: +async def test_stop_job(hass: HomeAssistant, init_integration: MockConfigEntry) -> None: """Test the stop job button.""" - coordinator: OctoprintDataUpdateCoordinator = hass.data[DOMAIN]["uuid"][ - "coordinator" - ] + coordinator: OctoprintDataUpdateCoordinator = init_integration.runtime_data # Test stopping the printer when it is paused with patch("pyoctoprintapi.OctoprintClient.cancel_job") as stop_command: