From 23831d746ea322e0d23abca705478963ed741146 Mon Sep 17 00:00:00 2001 From: epenet <6771947+epenet@users.noreply.github.com> Date: Tue, 29 Nov 2022 13:36:33 +0100 Subject: [PATCH] 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 --- .../components/scrape/config_flow.py | 47 +++++++++++++++- homeassistant/components/scrape/strings.json | 1 + .../components/scrape/translations/en.json | 1 + tests/components/scrape/test_config_flow.py | 54 ++++++++++++++++++- 4 files changed, 99 insertions(+), 4 deletions(-) diff --git a/homeassistant/components/scrape/config_flow.py b/homeassistant/components/scrape/config_flow.py index 8b598ca90be..9f683795e08 100644 --- a/homeassistant/components/scrape/config_flow.py +++ b/homeassistant/components/scrape/config_flow.py @@ -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, + ), } diff --git a/homeassistant/components/scrape/strings.json b/homeassistant/components/scrape/strings.json index d14b0916f6b..1ac50c695c6 100644 --- a/homeassistant/components/scrape/strings.json +++ b/homeassistant/components/scrape/strings.json @@ -54,6 +54,7 @@ "init": { "menu_options": { "add_sensor": "Add sensor", + "remove_sensor": "Remove sensor", "resource": "Configure resource" } }, diff --git a/homeassistant/components/scrape/translations/en.json b/homeassistant/components/scrape/translations/en.json index 868ed3abaa6..4fa261d77b3 100644 --- a/homeassistant/components/scrape/translations/en.json +++ b/homeassistant/components/scrape/translations/en.json @@ -60,6 +60,7 @@ "init": { "menu_options": { "add_sensor": "Add sensor", + "remove_sensor": "Remove sensor", "resource": "Configure resource" } }, diff --git a/tests/components/scrape/test_config_flow.py b/tests/components/scrape/test_config_flow.py index 9e7a895eadd..67a55f5a144 100644 --- a/tests/components/scrape/test_config_flow.py +++ b/tests/components/scrape/test_config_flow.py @@ -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"