Add ability to remove sensors in scrape config flow (#82912)

* Add ability to remove sensors in scrape config flow

* Add test

* Cleanup registry

* Cleanup

* Adjust comments

* Fix bad cleanup
This commit is contained in:
epenet 2022-11-29 13:36:33 +01:00 committed by GitHub
parent ce1b2f45c7
commit 23831d746e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 99 additions and 4 deletions

View File

@ -12,6 +12,7 @@ from homeassistant.components.rest.data import DEFAULT_TIMEOUT
from homeassistant.components.rest.schema import DEFAULT_METHOD, METHODS
from homeassistant.components.sensor import (
CONF_STATE_CLASS,
DOMAIN as SENSOR_DOMAIN,
SensorDeviceClass,
SensorStateClass,
)
@ -35,6 +36,7 @@ from homeassistant.const import (
UnitOfTemperature,
)
from homeassistant.core import async_get_hass
from homeassistant.helpers import config_validation as cv, entity_registry as er
from homeassistant.helpers.schema_config_entry_flow import (
SchemaCommonFlowHandler,
SchemaConfigFlowHandler,
@ -136,11 +138,47 @@ def validate_sensor_setup(
# Standard behavior is to merge the result with the options.
# In this case, we want to add a sub-item so we update the options directly.
sensors: list[dict[str, Any]] = handler.options.setdefault("sensor", [])
sensors: list[dict[str, Any]] = handler.options.setdefault(SENSOR_DOMAIN, [])
sensors.append(user_input)
return {}
def get_remove_sensor_schema(handler: SchemaCommonFlowHandler) -> vol.Schema:
"""Return schema for sensor removal."""
return vol.Schema(
{
vol.Required(CONF_INDEX): cv.multi_select(
{
str(index): config[CONF_NAME]
for index, config in enumerate(handler.options[SENSOR_DOMAIN])
},
)
}
)
def validate_remove_sensor(
handler: SchemaCommonFlowHandler, user_input: dict[str, Any]
) -> dict[str, Any]:
"""Validate remove sensor."""
removed_indexes: set[str] = set(user_input[CONF_INDEX])
# Standard behavior is to merge the result with the options.
# In this case, we want to remove sub-items so we update the options directly.
entity_registry = er.async_get(handler.parent_handler.hass)
sensors: list[dict[str, Any]] = []
sensor: dict[str, Any]
for index, sensor in enumerate(handler.options[SENSOR_DOMAIN]):
if str(index) not in removed_indexes:
sensors.append(sensor)
elif entity_id := entity_registry.async_get_entity_id(
SENSOR_DOMAIN, DOMAIN, sensor[CONF_UNIQUE_ID]
):
entity_registry.async_remove(entity_id)
handler.options[SENSOR_DOMAIN] = sensors
return {}
DATA_SCHEMA_RESOURCE = vol.Schema(RESOURCE_SETUP)
DATA_SCHEMA_SENSOR = vol.Schema(SENSOR_SETUP)
@ -156,7 +194,7 @@ CONFIG_FLOW = {
),
}
OPTIONS_FLOW = {
"init": SchemaFlowMenuStep(["resource", "add_sensor"]),
"init": SchemaFlowMenuStep(["resource", "add_sensor", "remove_sensor"]),
"resource": SchemaFlowFormStep(
DATA_SCHEMA_RESOURCE,
validate_user_input=validate_rest_setup,
@ -166,6 +204,11 @@ OPTIONS_FLOW = {
suggested_values=None,
validate_user_input=validate_sensor_setup,
),
"remove_sensor": SchemaFlowFormStep(
get_remove_sensor_schema,
suggested_values=None,
validate_user_input=validate_remove_sensor,
),
}

View File

@ -54,6 +54,7 @@
"init": {
"menu_options": {
"add_sensor": "Add sensor",
"remove_sensor": "Remove sensor",
"resource": "Configure resource"
}
},

View File

@ -60,6 +60,7 @@
"init": {
"menu_options": {
"add_sensor": "Add sensor",
"remove_sensor": "Remove sensor",
"resource": "Configure resource"
}
},

View File

@ -221,10 +221,10 @@ async def test_options_resource_flow(
assert state.state == "Hidden Version: 2021.12.10"
async def test_options_add_sensor_flow(
async def test_options_add_remove_sensor_flow(
hass: HomeAssistant, loaded_entry: MockConfigEntry
) -> None:
"""Test options flow to add a sensor."""
"""Test options flow to add and remove a sensor."""
state = hass.states.get("sensor.current_version")
assert state.state == "Current Version: 2021.12.10"
@ -290,3 +290,53 @@ async def test_options_add_sensor_flow(
state = hass.states.get("sensor.template")
assert state.state == "Trying to get"
# Now remove the original sensor
result = await hass.config_entries.options.async_init(loaded_entry.entry_id)
assert result["type"] == FlowResultType.MENU
assert result["step_id"] == "init"
result = await hass.config_entries.options.async_configure(
result["flow_id"],
{"next_step_id": "remove_sensor"},
)
assert result["type"] == FlowResultType.FORM
assert result["step_id"] == "remove_sensor"
mocker = MockRestData("test_scrape_sensor2")
with patch("homeassistant.components.rest.RestData", return_value=mocker):
result = await hass.config_entries.options.async_configure(
result["flow_id"],
user_input={
CONF_INDEX: ["0"],
},
)
await hass.async_block_till_done()
assert result["type"] == FlowResultType.CREATE_ENTRY
assert result["data"] == {
CONF_RESOURCE: "https://www.home-assistant.io",
CONF_METHOD: "GET",
CONF_VERIFY_SSL: True,
CONF_TIMEOUT: 10,
"sensor": [
{
CONF_NAME: "Template",
CONF_SELECT: "template",
CONF_INDEX: 0,
CONF_UNIQUE_ID: "3699ef88-69e6-11ed-a1eb-0242ac120003",
},
],
}
await hass.async_block_till_done()
# Check the original entity was removed, with only the new entity left
assert len(hass.states.async_all()) == 1
# Check the state of the new entity
state = hass.states.get("sensor.template")
assert state.state == "Trying to get"