From adf2fa5664d41a09c3f904e9df1343e08f8893b6 Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Fri, 3 Dec 2021 18:20:00 +0100
Subject: [PATCH] Use EntityCategory enum in Onewire (#60907)

* Use EntityCategory enum in Onewire

* Add checks for the entity_category

* Fix typo

Co-authored-by: epenet <epenet@users.noreply.github.com>
---
 homeassistant/components/onewire/switch.py | 5 +++--
 tests/components/onewire/__init__.py       | 4 ++++
 tests/components/onewire/const.py          | 3 +++
 3 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/homeassistant/components/onewire/switch.py b/homeassistant/components/onewire/switch.py
index 3146f4fb8a6..49f1ece51ca 100644
--- a/homeassistant/components/onewire/switch.py
+++ b/homeassistant/components/onewire/switch.py
@@ -9,8 +9,9 @@ from typing import TYPE_CHECKING, Any
 from homeassistant.components.onewire.model import OWServerDeviceDescription
 from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_TYPE, ENTITY_CATEGORY_CONFIG
+from homeassistant.const import CONF_TYPE
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import (
@@ -62,7 +63,7 @@ DEVICE_SWITCHES: dict[str, tuple[OneWireEntityDescription, ...]] = {
         OneWireSwitchEntityDescription(
             key="IAD",
             entity_registry_enabled_default=False,
-            entity_category=ENTITY_CATEGORY_CONFIG,
+            entity_category=EntityCategory.CONFIG,
             name="IAD",
             read_mode=READ_MODE_BOOL,
         ),
diff --git a/tests/components/onewire/__init__.py b/tests/components/onewire/__init__.py
index 36035c1c85b..3ae6dbab050 100644
--- a/tests/components/onewire/__init__.py
+++ b/tests/components/onewire/__init__.py
@@ -24,6 +24,7 @@ from homeassistant.helpers.entity_registry import EntityRegistry
 from .const import (
     ATTR_DEFAULT_DISABLED,
     ATTR_DEVICE_FILE,
+    ATTR_ENTITY_CATEGORY,
     ATTR_INJECT_READS,
     ATTR_UNIQUE_ID,
     FIXED_ATTRIBUTES,
@@ -77,6 +78,9 @@ def check_entities(
         entity_id = expected_entity[ATTR_ENTITY_ID]
         registry_entry = entity_registry.entities.get(entity_id)
         assert registry_entry is not None
+        assert registry_entry.entity_category == expected_entity.get(
+            ATTR_ENTITY_CATEGORY
+        )
         assert registry_entry.unique_id == expected_entity[ATTR_UNIQUE_ID]
         state = hass.states.get(entity_id)
         assert state.state == expected_entity[ATTR_STATE]
diff --git a/tests/components/onewire/const.py b/tests/components/onewire/const.py
index ebf52f2a25f..c38f5e53847 100644
--- a/tests/components/onewire/const.py
+++ b/tests/components/onewire/const.py
@@ -36,10 +36,12 @@ from homeassistant.const import (
     STATE_UNKNOWN,
     TEMP_CELSIUS,
 )
+from homeassistant.helpers.entity import EntityCategory
 
 ATTR_DEFAULT_DISABLED = "default_disabled"
 ATTR_DEVICE_FILE = "device_file"
 ATTR_DEVICE_INFO = "device_info"
+ATTR_ENTITY_CATEGORY = "entity_category"
 ATTR_INJECT_READS = "inject_reads"
 ATTR_UNIQUE_ID = "unique_id"
 ATTR_UNKNOWN_DEVICE = "unknown_device"
@@ -404,6 +406,7 @@ MOCK_OWPROXY_DEVICES = {
         SWITCH_DOMAIN: [
             {
                 ATTR_DEFAULT_DISABLED: True,
+                ATTR_ENTITY_CATEGORY: EntityCategory.CONFIG,
                 ATTR_ENTITY_ID: "switch.26_111111111111_iad",
                 ATTR_INJECT_READS: b"    1",
                 ATTR_STATE: STATE_ON,