From f36a300651819dd25f1a567e6ce1b8b05fba1bf3 Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Thu, 31 Aug 2023 15:16:32 +0200
Subject: [PATCH] Improve template sensor config flow validation (#99373)

---
 .../components/template/config_flow.py        | 27 +++++++---
 tests/components/template/test_config_flow.py | 54 +++++++++++++++++--
 2 files changed, 68 insertions(+), 13 deletions(-)

diff --git a/homeassistant/components/template/config_flow.py b/homeassistant/components/template/config_flow.py
index b89b3cbc91de..b2ccddedad8b 100644
--- a/homeassistant/components/template/config_flow.py
+++ b/homeassistant/components/template/config_flow.py
@@ -160,32 +160,43 @@ def _validate_unit(options: dict[str, Any]) -> None:
         and (units := DEVICE_CLASS_UNITS.get(device_class)) is not None
         and (unit := options.get(CONF_UNIT_OF_MEASUREMENT)) not in units
     ):
-        units_string = sorted(
-            [str(unit) if unit else "no unit of measurement" for unit in units],
+        sorted_units = sorted(
+            [f"'{str(unit)}'" if unit else "no unit of measurement" for unit in units],
             key=str.casefold,
         )
+        if len(sorted_units) == 1:
+            units_string = sorted_units[0]
+        else:
+            units_string = f"one of {', '.join(sorted_units)}"
 
         raise vol.Invalid(
             f"'{unit}' is not a valid unit for device class '{device_class}'; "
-            f"expected one of {', '.join(units_string)}"
+            f"expected {units_string}"
         )
 
 
 def _validate_state_class(options: dict[str, Any]) -> None:
     """Validate state class."""
     if (
-        (device_class := options.get(CONF_DEVICE_CLASS))
+        (state_class := options.get(CONF_STATE_CLASS))
+        and (device_class := options.get(CONF_DEVICE_CLASS))
         and (state_classes := DEVICE_CLASS_STATE_CLASSES.get(device_class)) is not None
-        and (state_class := options.get(CONF_STATE_CLASS)) not in state_classes
+        and state_class not in state_classes
     ):
-        state_classes_string = sorted(
-            [str(state_class) for state_class in state_classes],
+        sorted_state_classes = sorted(
+            [f"'{str(state_class)}'" for state_class in state_classes],
             key=str.casefold,
         )
+        if len(sorted_state_classes) == 0:
+            state_classes_string = "no state class"
+        elif len(sorted_state_classes) == 1:
+            state_classes_string = sorted_state_classes[0]
+        else:
+            state_classes_string = f"one of {', '.join(sorted_state_classes)}"
 
         raise vol.Invalid(
             f"'{state_class}' is not a valid state class for device class "
-            f"'{device_class}'; expected one of {', '.join(state_classes_string)}"
+            f"'{device_class}'; expected {state_classes_string}"
         )
 
 
diff --git a/tests/components/template/test_config_flow.py b/tests/components/template/test_config_flow.py
index dd283ff92144..ba939f3b8d10 100644
--- a/tests/components/template/test_config_flow.py
+++ b/tests/components/template/test_config_flow.py
@@ -349,18 +349,62 @@ EARLY_END_ERROR = "invalid template (TemplateSyntaxError: unexpected 'end of tem
     [
         ("binary_sensor", "{{", {}, {"state": EARLY_END_ERROR}),
         ("sensor", "{{", {}, {"state": EARLY_END_ERROR}),
+        (
+            "sensor",
+            "",
+            {"device_class": "aqi", "unit_of_measurement": "cats"},
+            {
+                "unit_of_measurement": (
+                    "'cats' is not a valid unit for device class 'aqi'; "
+                    "expected no unit of measurement"
+                ),
+            },
+        ),
         (
             "sensor",
             "",
             {"device_class": "temperature", "unit_of_measurement": "cats"},
             {
-                "state_class": (
-                    "'None' is not a valid state class for device class 'temperature'; "
-                    "expected one of measurement"
-                ),
                 "unit_of_measurement": (
                     "'cats' is not a valid unit for device class 'temperature'; "
-                    "expected one of K, °C, °F"
+                    "expected one of 'K', '°C', '°F'"
+                ),
+            },
+        ),
+        (
+            "sensor",
+            "",
+            {"device_class": "timestamp", "state_class": "measurement"},
+            {
+                "state_class": (
+                    "'measurement' is not a valid state class for device class "
+                    "'timestamp'; expected no state class"
+                ),
+            },
+        ),
+        (
+            "sensor",
+            "",
+            {"device_class": "aqi", "state_class": "total"},
+            {
+                "state_class": (
+                    "'total' is not a valid state class for device class "
+                    "'aqi'; expected 'measurement'"
+                ),
+            },
+        ),
+        (
+            "sensor",
+            "",
+            {"device_class": "energy", "state_class": "measurement"},
+            {
+                "state_class": (
+                    "'measurement' is not a valid state class for device class "
+                    "'energy'; expected one of 'total', 'total_increasing'"
+                ),
+                "unit_of_measurement": (
+                    "'None' is not a valid unit for device class 'energy'; "
+                    "expected one of 'GJ', 'kWh', 'MJ', 'MWh', 'Wh'"
                 ),
             },
         ),