diff --git a/homeassistant/components/nextbus/__init__.py b/homeassistant/components/nextbus/__init__.py index 5d0029f4448e..02d5a8cd1e2f 100644 --- a/homeassistant/components/nextbus/__init__.py +++ b/homeassistant/components/nextbus/__init__.py @@ -4,31 +4,33 @@ from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_STOP, Platform from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryNotReady +from homeassistant.util.hass_dict import HassKey from .const import CONF_AGENCY, CONF_ROUTE, DOMAIN from .coordinator import NextBusDataUpdateCoordinator PLATFORMS = [Platform.SENSOR] +type NextBusConfigEntry = ConfigEntry[NextBusDataUpdateCoordinator] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +# Coordinators are shared between entries with the same agency and stop; the +# synchronous check and store below must stay free of awaits so concurrent +# entry setups cannot create duplicates. +NEXTBUS_KEY: HassKey[dict[str, NextBusDataUpdateCoordinator]] = HassKey(DOMAIN) + + +async def async_setup_entry(hass: HomeAssistant, entry: NextBusConfigEntry) -> bool: """Set up platforms for NextBus.""" entry_agency = entry.data[CONF_AGENCY] entry_stop = entry.data[CONF_STOP] coordinator_key = f"{entry_agency}-{entry_stop}" - # Uses legacy hass.data[DOMAIN] pattern - # pylint: disable-next=home-assistant-use-runtime-data - coordinator: NextBusDataUpdateCoordinator | None = hass.data.setdefault( - DOMAIN, {} - ).get( - coordinator_key, - ) + coordinators = hass.data.setdefault(NEXTBUS_KEY, {}) + coordinator = coordinators.get(coordinator_key) if coordinator is None: coordinator = NextBusDataUpdateCoordinator(hass, entry_agency) - # Uses legacy hass.data[DOMAIN] pattern - # pylint: disable-next=home-assistant-use-runtime-data - hass.data[DOMAIN][coordinator_key] = coordinator + coordinators[coordinator_key] = coordinator + entry.runtime_data = coordinator coordinator.add_stop_route(entry_stop, entry.data[CONF_ROUTE]) @@ -41,19 +43,17 @@ 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: NextBusConfigEntry) -> bool: """Unload a config entry.""" if await hass.config_entries.async_unload_platforms(entry, PLATFORMS): entry_agency = entry.data[CONF_AGENCY] entry_stop = entry.data[CONF_STOP] - coordinator_key = f"{entry_agency}-{entry_stop}" - - coordinator: NextBusDataUpdateCoordinator = hass.data[DOMAIN][coordinator_key] + coordinator = entry.runtime_data coordinator.remove_stop_route(entry_stop, entry.data[CONF_ROUTE]) if not coordinator.has_routes(): await coordinator.async_shutdown() - hass.data[DOMAIN].pop(coordinator_key) + hass.data[NEXTBUS_KEY].pop(f"{entry_agency}-{entry_stop}") return True diff --git a/homeassistant/components/nextbus/sensor.py b/homeassistant/components/nextbus/sensor.py index 8d09055a7b0b..3160b1101037 100644 --- a/homeassistant/components/nextbus/sensor.py +++ b/homeassistant/components/nextbus/sensor.py @@ -4,14 +4,14 @@ import logging from typing import cast, override from homeassistant.components.sensor import SensorDeviceClass, SensorEntity -from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_NAME, CONF_STOP from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback from homeassistant.helpers.update_coordinator import CoordinatorEntity from homeassistant.util.dt import utc_from_timestamp -from .const import CONF_AGENCY, CONF_ROUTE, DOMAIN +from . import NextBusConfigEntry +from .const import CONF_AGENCY, CONF_ROUTE from .coordinator import NextBusDataUpdateCoordinator from .util import maybe_first @@ -20,18 +20,12 @@ _LOGGER = logging.getLogger(__name__) async def async_setup_entry( hass: HomeAssistant, - config: ConfigEntry, + config: NextBusConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Load values from configuration and initialize the platform.""" _LOGGER.debug(config.data) - entry_agency = config.data[CONF_AGENCY] - entry_stop = config.data[CONF_STOP] - coordinator_key = f"{entry_agency}-{entry_stop}" - - # Uses legacy hass.data[DOMAIN] pattern - # pylint: disable-next=home-assistant-use-runtime-data - coordinator: NextBusDataUpdateCoordinator = hass.data[DOMAIN].get(coordinator_key) + coordinator = config.runtime_data async_add_entities( ( diff --git a/tests/components/nextbus/const.py b/tests/components/nextbus/const.py index 66eb3635ca92..b76af217adf0 100644 --- a/tests/components/nextbus/const.py +++ b/tests/components/nextbus/const.py @@ -7,7 +7,6 @@ from homeassistant.const import CONF_STOP VALID_AGENCY = "sfmta-cis" VALID_ROUTE = "F" VALID_STOP = "5184" -VALID_COORDINATOR_KEY = f"{VALID_AGENCY}-{VALID_STOP}" VALID_AGENCY_TITLE = "San Francisco Muni" VALID_ROUTE_TITLE = "F-Market & Wharves" VALID_STOP_TITLE = "Market St & 7th St" diff --git a/tests/components/nextbus/test_sensor.py b/tests/components/nextbus/test_sensor.py index faba9c7c8f9d..0f84b7c9df8d 100644 --- a/tests/components/nextbus/test_sensor.py +++ b/tests/components/nextbus/test_sensor.py @@ -9,12 +9,14 @@ from freezegun.api import FrozenDateTimeFactory from py_nextbus.client import NextBusFormatError, NextBusHTTPError import pytest -from homeassistant.components.nextbus.const import DOMAIN +from homeassistant.components.nextbus import NEXTBUS_KEY +from homeassistant.components.nextbus.const import CONF_AGENCY, CONF_ROUTE, DOMAIN from homeassistant.components.nextbus.coordinator import NextBusDataUpdateCoordinator from homeassistant.config_entries import ConfigEntryState -from homeassistant.const import CONF_NAME +from homeassistant.const import CONF_NAME, CONF_STOP from homeassistant.core import HomeAssistant from homeassistant.helpers.update_coordinator import UpdateFailed +from homeassistant.setup import async_setup_component from homeassistant.util import dt as dt_util from . import assert_setup_sensor @@ -27,12 +29,12 @@ from .const import ( SENSOR_ID, SENSOR_ID_2, VALID_AGENCY, - VALID_COORDINATOR_KEY, + VALID_AGENCY_TITLE, VALID_ROUTE_TITLE, VALID_STOP_TITLE, ) -from tests.common import async_fire_time_changed +from tests.common import MockConfigEntry, async_fire_time_changed async def test_predictions( @@ -69,8 +71,8 @@ async def test_prediction_exceptions( client_exception: Exception, ) -> None: """Test that some coodinator exceptions raise UpdateFailed exceptions.""" - await assert_setup_sensor(hass, CONFIG_BASIC) - coordinator: NextBusDataUpdateCoordinator = hass.data[DOMAIN][VALID_COORDINATOR_KEY] + entry = await assert_setup_sensor(hass, CONFIG_BASIC) + coordinator: NextBusDataUpdateCoordinator = entry.runtime_data mock_nextbus_predictions.side_effect = client_exception with pytest.raises(UpdateFailed): await coordinator._async_update_data() @@ -175,6 +177,39 @@ async def test_verify_throttle( assert state.state == "unknown" +async def test_concurrent_setup_shares_coordinator( + hass: HomeAssistant, + mock_nextbus: MagicMock, + mock_nextbus_lists: MagicMock, + mock_nextbus_predictions: MagicMock, +) -> None: + """Test that two entries set up concurrently share one coordinator.""" + entries = [] + for config, route_title in ( + (CONFIG_BASIC, VALID_ROUTE_TITLE), + (CONFIG_BASIC_2, ROUTE_TITLE_2), + ): + entry = MockConfigEntry( + domain=DOMAIN, + data=config[DOMAIN], + title=f"{VALID_AGENCY_TITLE} {route_title} {VALID_STOP_TITLE}", + unique_id=( + f"{config[DOMAIN][CONF_AGENCY]}" + f"_{config[DOMAIN][CONF_ROUTE]}" + f"_{config[DOMAIN][CONF_STOP]}" + ), + ) + entry.add_to_hass(hass) + entries.append(entry) + + assert await async_setup_component(hass, DOMAIN, {}) + await hass.async_block_till_done() + + assert entries[0].state is ConfigEntryState.LOADED + assert entries[1].state is ConfigEntryState.LOADED + assert entries[0].runtime_data is entries[1].runtime_data + + async def test_unload_entry( hass: HomeAssistant, mock_nextbus: MagicMock, @@ -184,7 +219,11 @@ async def test_unload_entry( ) -> None: """Test that the sensor can be unloaded.""" config_entry1 = await assert_setup_sensor(hass, CONFIG_BASIC) - await assert_setup_sensor(hass, CONFIG_BASIC_2, route_title=ROUTE_TITLE_2) + config_entry2 = await assert_setup_sensor( + hass, CONFIG_BASIC_2, route_title=ROUTE_TITLE_2 + ) + + assert config_entry1.runtime_data is config_entry2.runtime_data # Verify the first sensor state = hass.states.get(SENSOR_ID) @@ -224,3 +263,27 @@ async def test_unload_entry( assert state is not None assert state.attributes["upcoming"] == "5" assert state.state == "2019-03-28T21:09:35+00:00" + + +async def test_unload_final_entry_cleans_up_shared_coordinator( + hass: HomeAssistant, + mock_nextbus: MagicMock, + mock_nextbus_lists: MagicMock, + mock_nextbus_predictions: MagicMock, +) -> None: + """Test that unloading the final entry shuts down the shared coordinator.""" + config_entry1 = await assert_setup_sensor(hass, CONFIG_BASIC) + config_entry2 = await assert_setup_sensor( + hass, CONFIG_BASIC_2, route_title=ROUTE_TITLE_2 + ) + coordinator: NextBusDataUpdateCoordinator = config_entry1.runtime_data + + await hass.config_entries.async_unload(config_entry1.entry_id) + await hass.async_block_till_done() + await hass.config_entries.async_unload(config_entry2.entry_id) + await hass.async_block_till_done() + + assert config_entry1.state is ConfigEntryState.NOT_LOADED + assert config_entry2.state is ConfigEntryState.NOT_LOADED + assert coordinator._shutdown_requested + assert hass.data[NEXTBUS_KEY] == {}