Add all kraken entities on startup (#58027)

* Add all kraken entities on startup

* Add _async_add_kraken_sensors

* Add test sensors available after restart

* Add test for config update

* Update tests/components/kraken/test_sensor.py

Co-authored-by: Allen Porter <allen.porter@gmail.com>

Co-authored-by: Allen Porter <allen.porter@gmail.com>
pull/58531/head
Kevin Stillhammer 2021-10-26 09:02:49 +02:00 committed by GitHub
parent cee51ead91
commit bd5c131675
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 107 additions and 14 deletions

View File

@ -33,8 +33,26 @@ async def async_setup_entry(
) -> None:
"""Add kraken entities from a config_entry."""
def _async_add_kraken_sensors(tracked_asset_pairs: list[str]) -> None:
entities = []
for tracked_asset_pair in tracked_asset_pairs:
entities.extend(
[
KrakenSensor(
hass.data[DOMAIN],
tracked_asset_pair,
description,
)
for description in SENSOR_TYPES
]
)
async_add_entities(entities, True)
_async_add_kraken_sensors(config_entry.options[CONF_TRACKED_ASSET_PAIRS])
@callback
def async_update_sensors(hass: HomeAssistant, config_entry: ConfigEntry) -> None:
"""Add or remove sensors for configured tracked asset pairs."""
dev_reg = device_registry.async_get(hass)
existing_devices = {
@ -44,7 +62,7 @@ async def async_setup_entry(
)
}
entities = []
asset_pairs_to_add = []
for tracked_asset_pair in config_entry.options[CONF_TRACKED_ASSET_PAIRS]:
# Only create new devices
if (
@ -52,24 +70,13 @@ async def async_setup_entry(
) in existing_devices:
existing_devices.pop(device_name)
else:
entities.extend(
[
KrakenSensor(
hass.data[DOMAIN],
tracked_asset_pair,
description,
)
for description in SENSOR_TYPES
]
)
async_add_entities(entities, True)
asset_pairs_to_add.append(tracked_asset_pair)
_async_add_kraken_sensors(asset_pairs_to_add)
# Remove devices for asset pairs which are no longer tracked
for device_id in existing_devices.values():
dev_reg.async_remove_device(device_id)
async_update_sensors(hass, config_entry)
config_entry.async_on_unload(
async_dispatcher_connect(
hass,

View File

@ -228,6 +228,92 @@ async def test_sensor(hass):
assert xbt_usd_opening_price_today.state == "0.0003513"
async def test_sensors_available_after_restart(hass):
"""Test that all sensors are added again after a restart."""
utcnow = dt_util.utcnow()
# Patching 'utcnow' to gain more control over the timed update.
with patch("homeassistant.util.dt.utcnow", return_value=utcnow), patch(
"pykrakenapi.KrakenAPI.get_tradable_asset_pairs",
return_value=TRADEABLE_ASSET_PAIR_RESPONSE,
), patch(
"pykrakenapi.KrakenAPI.get_ticker_information",
return_value=TICKER_INFORMATION_RESPONSE,
):
entry = MockConfigEntry(
domain=DOMAIN,
options={
CONF_SCAN_INTERVAL: DEFAULT_SCAN_INTERVAL,
CONF_TRACKED_ASSET_PAIRS: [DEFAULT_TRACKED_ASSET_PAIR],
},
)
device_registry = await hass.helpers.device_registry.async_get_registry()
device_registry.async_get_or_create(
config_entry_id=entry.entry_id,
identifiers={(DOMAIN, "XBT_USD")},
name="XBT USD",
manufacturer="Kraken.com",
entry_type="service",
)
entry.add_to_hass(hass)
await hass.config_entries.async_setup(entry.entry_id)
await hass.async_block_till_done()
hass.bus.async_fire(EVENT_HOMEASSISTANT_START)
await hass.async_block_till_done()
sensor = hass.states.get("sensor.xbt_usd_ask")
assert sensor.state == "0.0003494"
async def test_sensors_added_after_config_update(hass):
"""Test that sensors are added when another tracked asset pair is added."""
utcnow = dt_util.utcnow()
# Patching 'utcnow' to gain more control over the timed update.
with patch("homeassistant.util.dt.utcnow", return_value=utcnow), patch(
"pykrakenapi.KrakenAPI.get_tradable_asset_pairs",
return_value=TRADEABLE_ASSET_PAIR_RESPONSE,
), patch(
"pykrakenapi.KrakenAPI.get_ticker_information",
return_value=TICKER_INFORMATION_RESPONSE,
):
entry = MockConfigEntry(
domain=DOMAIN,
options={
CONF_SCAN_INTERVAL: DEFAULT_SCAN_INTERVAL,
CONF_TRACKED_ASSET_PAIRS: [DEFAULT_TRACKED_ASSET_PAIR],
},
)
entry.add_to_hass(hass)
await hass.config_entries.async_setup(entry.entry_id)
await hass.async_block_till_done()
hass.bus.async_fire(EVENT_HOMEASSISTANT_START)
await hass.async_block_till_done()
assert hass.states.get("sensor.xbt_usd_ask")
assert not hass.states.get("sensor.ada_xbt_ask")
hass.config_entries.async_update_entry(
entry,
options={
CONF_SCAN_INTERVAL: DEFAULT_SCAN_INTERVAL,
CONF_TRACKED_ASSET_PAIRS: [DEFAULT_TRACKED_ASSET_PAIR, "ADA/XBT"],
},
)
async_fire_time_changed(
hass, utcnow + timedelta(seconds=DEFAULT_SCAN_INTERVAL * 2)
)
await hass.async_block_till_done()
assert hass.states.get("sensor.ada_xbt_ask")
async def test_missing_pair_marks_sensor_unavailable(hass):
"""Test that a missing tradable asset pair marks the sensor unavailable."""
utcnow = dt_util.utcnow()