diff --git a/esphome/core/entity_helpers.py b/esphome/core/entity_helpers.py
index cff6e6bd4b..a3d6a3b2fe 100644
--- a/esphome/core/entity_helpers.py
+++ b/esphome/core/entity_helpers.py
@@ -228,9 +228,6 @@ def setup_unit_of_measurement(config: ConfigType) -> None:
     config[_KEY_UOM_IDX] = idx
 
 
-_ENTITY_CATEGORY_NAMES = {0: "", 1: "config", 2: "diagnostic"}
-
-
 def _sanitize_comment(text: str) -> str:
     r"""Sanitize a string for safe inclusion in a C++ // line comment.
 
@@ -248,7 +245,10 @@ def _describe_packed_flags(config: ConfigType, entity_category: int) -> str:
         parts.append("internal")
     if config.get(_KEY_DISABLED_BY_DEFAULT):
         parts.append("disabled_by_default")
-    if cat_name := _ENTITY_CATEGORY_NAMES.get(entity_category, ""):
+    entity_cat_keys = list(cv.ENTITY_CATEGORIES)
+    if entity_category < len(entity_cat_keys) and (
+        cat_name := entity_cat_keys[entity_category]
+    ):
         parts.append(f"category:{cat_name}")
     if dc := config.get(CONF_DEVICE_CLASS):
         parts.append(f"dc:{_sanitize_comment(dc)}")
diff --git a/tests/unit_tests/core/test_entity_helpers.py b/tests/unit_tests/core/test_entity_helpers.py
index 3a988ec274..4e24a21b68 100644
--- a/tests/unit_tests/core/test_entity_helpers.py
+++ b/tests/unit_tests/core/test_entity_helpers.py
@@ -942,7 +942,6 @@ async def test_setup_entity_with_entity_category(
     setup_test_environment: list[str],
 ) -> None:
     """Test setup_entity sets entity_category correctly."""
-    setup_test_environment  # noqa: F841 - fixture initializes CORE state
     var = MockObj("sensor1")
     config = {
         CONF_NAME: "Temperature",
@@ -1008,8 +1007,6 @@ async def test_setup_entity_decorator_mode(setup_test_environment: list[str]) ->
 
 def _extract_packed_value(expressions: list[str]) -> int:
     """Extract the third argument (packed value) from a configure_entity_() call."""
-    import re
-
     for expr in expressions:
         if "configure_entity_" in expr:
             # Match the last integer argument before the closing ");"