329 files

custom_components/powercalc/flow_helper/flows/cost.py

@@ -1,12 +1,11 @@
 """Config/options flow for a standalone cost sensor based on an existing energy sensor."""
 
-from __future__ import annotations
-
 from typing import TYPE_CHECKING, Any
 
 from homeassistant.components.sensor import SensorDeviceClass
 from homeassistant.config_entries import ConfigFlowResult
 from homeassistant.const import CONF_NAME
+from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import section
 from homeassistant.helpers import selector
 from homeassistant.helpers.schema_config_entry_flow import SchemaFlowError
@@ -21,28 +20,35 @@ from custom_components.powercalc.const import (
     SensorType,
 )
 from custom_components.powercalc.flow_helper.common import PowercalcFormStep, Step, flatten_sections
-from custom_components.powercalc.flow_helper.schema import SCHEMA_COST_PRICING, SECTION_COST_PRICING
+from custom_components.powercalc.flow_helper.schema import SECTION_COST_PRICING, build_cost_pricing_schema
 
 if TYPE_CHECKING:
-    from homeassistant.core import HomeAssistant
-
     from custom_components.powercalc.config_flow import PowercalcConfigFlow, PowercalcOptionsFlow
 
-SCHEMA_COST_OPTIONS = vol.Schema(
-    {
-        vol.Required(CONF_ENERGY_SENSOR_ID): selector.EntitySelector(
-            selector.EntitySelectorConfig(domain="sensor", device_class=SensorDeviceClass.ENERGY),
-        ),
-        vol.Optional(SECTION_COST_PRICING): section(SCHEMA_COST_PRICING, {"collapsed": True}),
-    },
-)
 
-SCHEMA_COST = vol.Schema(
-    {
-        vol.Required(CONF_NAME): selector.TextSelector(),
-        **SCHEMA_COST_OPTIONS.schema,
-    },
-)
+def build_cost_options_schema(hass: HomeAssistant, current_price_entity_id: str | None = None) -> vol.Schema:
+    """Return the schema for the options of a standalone cost sensor."""
+    return vol.Schema(
+        {
+            vol.Required(CONF_ENERGY_SENSOR_ID): selector.EntitySelector(
+                selector.EntitySelectorConfig(domain="sensor", device_class=SensorDeviceClass.ENERGY),
+            ),
+            vol.Optional(SECTION_COST_PRICING): section(
+                build_cost_pricing_schema(hass, current_price_entity_id),
+                {"collapsed": True},
+            ),
+        },
+    )
+
+
+def build_cost_schema(hass: HomeAssistant) -> vol.Schema:
+    """Return the schema for creating a standalone cost sensor."""
+    return vol.Schema(
+        {
+            vol.Required(CONF_NAME): selector.TextSelector(),
+            **build_cost_options_schema(hass).schema,
+        },
+    )
 
 
 def is_global_price_configured(hass: HomeAssistant) -> bool:
@@ -68,13 +74,17 @@ class CostConfigFlow:
         """Handle the flow for a standalone cost sensor."""
         self.flow.selected_sensor_type = SensorType.COST
         return await self.flow.handle_form_step(
-            PowercalcFormStep(step=Step.COST, schema=SCHEMA_COST, validate_user_input=self.validate),
+            PowercalcFormStep(
+                step=Step.COST,
+                schema=build_cost_schema(self.flow.hass),
+                validate_user_input=self.validate,
+            ),
             user_input,
         )
 
     def validate(self, user_input: dict[str, Any]) -> dict[str, Any]:
         """Flatten the pricing section and reject the input when no price is available."""
-        user_input = flatten_sections(user_input, SCHEMA_COST)
+        user_input = flatten_sections(user_input, build_cost_schema(self.flow.hass))
         if validate_price_configured(self.flow.hass, user_input):
             raise SchemaFlowError("cost_price_mandatory")
         return user_input
@@ -88,7 +98,7 @@ class CostOptionsFlow:
         """Handle the cost sensor options flow."""
         return await self.flow.async_handle_options_step(
             user_input,
-            SCHEMA_COST_OPTIONS,
+            build_cost_options_schema(self.flow.hass, self.flow.sensor_config.get(CONF_ENERGY_PRICE_SENSOR)),
             Step.COST,
             validate=lambda flat_input: validate_price_configured(self.flow.hass, flat_input),
         )