Add optional limits to compensation sensors (#85886)

Co-authored-by: Tom Harris <tomharris@harrisnj.net>
Co-authored-by: J. Nick Koston <nick@koston.org>

homeassistant/components/compensation/__init__.py

@@ -1,5 +1,6 @@
 """The Compensation integration."""
 import logging
+from operator import itemgetter
 
 import numpy as np
 import voluptuous as vol
@@ -7,6 +8,8 @@ import voluptuous as vol
 from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
 from homeassistant.const import (
     CONF_ATTRIBUTE,
+    CONF_MAXIMUM,
+    CONF_MINIMUM,
     CONF_SOURCE,
     CONF_UNIQUE_ID,
     CONF_UNIT_OF_MEASUREMENT,
@@ -20,8 +23,10 @@ from .const import (
     CONF_COMPENSATION,
     CONF_DATAPOINTS,
     CONF_DEGREE,
+    CONF_LOWER_LIMIT,
     CONF_POLYNOMIAL,
     CONF_PRECISION,
+    CONF_UPPER_LIMIT,
     DATA_COMPENSATION,
     DEFAULT_DEGREE,
     DEFAULT_PRECISION,
@@ -50,6 +55,8 @@ COMPENSATION_SCHEMA = vol.Schema(
         ],
         vol.Optional(CONF_UNIQUE_ID): cv.string,
         vol.Optional(CONF_ATTRIBUTE): cv.string,
+        vol.Optional(CONF_UPPER_LIMIT, default=False): cv.boolean,
+        vol.Optional(CONF_LOWER_LIMIT, default=False): cv.boolean,
         vol.Optional(CONF_PRECISION, default=DEFAULT_PRECISION): cv.positive_int,
         vol.Optional(CONF_DEGREE, default=DEFAULT_DEGREE): vol.All(
             vol.Coerce(int),
@@ -78,8 +85,11 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
 
         degree = conf[CONF_DEGREE]
 
+        initial_coefficients: list[tuple[float, float]] = conf[CONF_DATAPOINTS]
+        sorted_coefficients = sorted(initial_coefficients, key=itemgetter(0))
+
         # get x values and y values from the x,y point pairs
-        x_values, y_values = zip(*conf[CONF_DATAPOINTS])
+        x_values, y_values = zip(*initial_coefficients)
 
         # try to get valid coefficients for a polynomial
         coefficients = None
@@ -99,6 +109,16 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
             }
             data[CONF_POLYNOMIAL] = np.poly1d(coefficients)
 
+            if data[CONF_LOWER_LIMIT]:
+                data[CONF_MINIMUM] = sorted_coefficients[0]
+            else:
+                data[CONF_MINIMUM] = None
+
+            if data[CONF_UPPER_LIMIT]:
+                data[CONF_MAXIMUM] = sorted_coefficients[-1]
+            else:
+                data[CONF_MAXIMUM] = None
+
             hass.data[DATA_COMPENSATION][compensation] = data
 
             hass.async_create_task(

homeassistant/components/compensation/const.py

@@ -4,6 +4,8 @@ DOMAIN = "compensation"
 SENSOR = "compensation"
 
 CONF_COMPENSATION = "compensation"
+CONF_LOWER_LIMIT = "lower_limit"
+CONF_UPPER_LIMIT = "upper_limit"
 CONF_DATAPOINTS = "data_points"
 CONF_DEGREE = "degree"
 CONF_PRECISION = "precision"

homeassistant/components/compensation/sensor.py

@@ -10,6 +10,8 @@ from homeassistant.components.sensor import SensorEntity
 from homeassistant.const import (
     ATTR_UNIT_OF_MEASUREMENT,
     CONF_ATTRIBUTE,
+    CONF_MAXIMUM,
+    CONF_MINIMUM,
     CONF_SOURCE,
     CONF_UNIQUE_ID,
     CONF_UNIT_OF_MEASUREMENT,
@@ -64,6 +66,8 @@ async def async_setup_platform(
                 conf[CONF_PRECISION],
                 conf[CONF_POLYNOMIAL],
                 conf.get(CONF_UNIT_OF_MEASUREMENT),
+                conf[CONF_MINIMUM],
+                conf[CONF_MAXIMUM],
             )
         ]
     )
@@ -83,6 +87,8 @@ class CompensationSensor(SensorEntity):
         precision: int,
         polynomial: np.poly1d,
         unit_of_measurement: str | None,
+        minimum: tuple[float, float] | None,
+        maximum: tuple[float, float] | None,
     ) -> None:
         """Initialize the Compensation sensor."""
         self._source_entity_id = source
@@ -93,6 +99,8 @@ class CompensationSensor(SensorEntity):
         self._coefficients = polynomial.coefficients.tolist()
         self._attr_unique_id = unique_id
         self._attr_name = name
+        self._minimum = minimum
+        self._maximum = maximum
 
     async def async_added_to_hass(self) -> None:
         """Handle added to Hass."""
@@ -132,7 +140,14 @@ class CompensationSensor(SensorEntity):
         else:
             value = None if new_state.state == STATE_UNKNOWN else new_state.state
         try:
-            self._attr_native_value = round(self._poly(float(value)), self._precision)
+            x_value = float(value)
+            if self._minimum is not None and x_value <= self._minimum[0]:
+                y_value = self._minimum[1]
+            elif self._maximum is not None and x_value >= self._maximum[0]:
+                y_value = self._maximum[1]
+            else:
+                y_value = self._poly(x_value)
+            self._attr_native_value = round(y_value, self._precision)
 
         except (ValueError, TypeError):
             self._attr_native_value = None

tests/components/compensation/test_sensor.py

@@ -223,3 +223,50 @@ async def test_new_state_is_none(hass: HomeAssistant) -> None:
     )
 
     assert last_changed == hass.states.get(expected_entity_id).last_changed
+
+
+@pytest.mark.parametrize(
+    ("lower", "upper"),
+    [
+        (True, False),
+        (False, True),
+        (True, True),
+    ],
+)
+async def test_limits(hass: HomeAssistant, lower: bool, upper: bool) -> None:
+    """Test compensation sensor state."""
+    source = "sensor.test"
+    config = {
+        "compensation": {
+            "test": {
+                "source": source,
+                "data_points": [
+                    [1.0, 0.0],
+                    [3.0, 2.0],
+                    [2.0, 1.0],
+                ],
+                "precision": 2,
+                "lower_limit": lower,
+                "upper_limit": upper,
+                "unit_of_measurement": "a",
+            }
+        }
+    }
+    await async_setup_component(hass, DOMAIN, config)
+    await hass.async_block_till_done()
+    await hass.async_start()
+    await hass.async_block_till_done()
+
+    entity_id = "sensor.compensation_sensor_test"
+
+    hass.states.async_set(source, 0, {})
+    await hass.async_block_till_done()
+    state = hass.states.get(entity_id)
+    value = 0.0 if lower else -1.0
+    assert float(state.state) == value
+
+    hass.states.async_set(source, 5, {})
+    await hass.async_block_till_done()
+    state = hass.states.get(entity_id)
+    value = 2.0 if upper else 4.0
+    assert float(state.state) == value