Use runtime_data in octoprint integration (#167028)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

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,

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

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
 

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

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")

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

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

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: