From 8a90ce882a373d9f682231dd23b1bb31438de340 Mon Sep 17 00:00:00 2001
From: Jesse Hills <3060199+jesserockz@users.noreply.github.com>
Date: Fri, 9 May 2025 18:22:43 +1200
Subject: [PATCH] [update] Move to use ``update_schema(..)`` instead of
 ``UPDATE_SCHEMA`` (#8726)

---
 .../http_request/update/__init__.py           | 19 +++++++------
 esphome/components/update/__init__.py         | 27 ++++++++++++++++++-
 2 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/esphome/components/http_request/update/__init__.py b/esphome/components/http_request/update/__init__.py
index f1b282d891..abb4b2a430 100644
--- a/esphome/components/http_request/update/__init__.py
+++ b/esphome/components/http_request/update/__init__.py
@@ -16,14 +16,17 @@ HttpRequestUpdate = http_request_ns.class_(
 
 CONF_OTA_ID = "ota_id"
 
-CONFIG_SCHEMA = update.UPDATE_SCHEMA.extend(
-    {
-        cv.GenerateID(): cv.declare_id(HttpRequestUpdate),
-        cv.GenerateID(CONF_OTA_ID): cv.use_id(OtaHttpRequestComponent),
-        cv.GenerateID(CONF_HTTP_REQUEST_ID): cv.use_id(HttpRequestComponent),
-        cv.Required(CONF_SOURCE): cv.url,
-    }
-).extend(cv.polling_component_schema("6h"))
+CONFIG_SCHEMA = (
+    update.update_schema(HttpRequestUpdate)
+    .extend(
+        {
+            cv.GenerateID(CONF_OTA_ID): cv.use_id(OtaHttpRequestComponent),
+            cv.GenerateID(CONF_HTTP_REQUEST_ID): cv.use_id(HttpRequestComponent),
+            cv.Required(CONF_SOURCE): cv.url,
+        }
+    )
+    .extend(cv.polling_component_schema("6h"))
+)
 
 
 async def to_code(config):
diff --git a/esphome/components/update/__init__.py b/esphome/components/update/__init__.py
index 4729d954ee..a607aefea0 100644
--- a/esphome/components/update/__init__.py
+++ b/esphome/components/update/__init__.py
@@ -6,6 +6,7 @@ from esphome.const import (
     CONF_DEVICE_CLASS,
     CONF_ENTITY_CATEGORY,
     CONF_FORCE_UPDATE,
+    CONF_ICON,
     CONF_ID,
     CONF_MQTT_ID,
     CONF_WEB_SERVER,
@@ -14,6 +15,7 @@ from esphome.const import (
     ENTITY_CATEGORY_CONFIG,
 )
 from esphome.core import CORE, coroutine_with_priority
+from esphome.cpp_generator import MockObjClass
 from esphome.cpp_helpers import setup_entity
 
 CODEOWNERS = ["@jesserockz"]
@@ -38,7 +40,7 @@ DEVICE_CLASSES = [
 
 CONF_ON_UPDATE_AVAILABLE = "on_update_available"
 
-UPDATE_SCHEMA = (
+_UPDATE_SCHEMA = (
     cv.ENTITY_BASE_SCHEMA.extend(web_server.WEBSERVER_SORTING_SCHEMA)
     .extend(cv.MQTT_COMMAND_COMPONENT_SCHEMA)
     .extend(
@@ -56,6 +58,29 @@ UPDATE_SCHEMA = (
 )
 
 
+def update_schema(
+    class_: MockObjClass = cv.UNDEFINED,
+    *,
+    icon: str = cv.UNDEFINED,
+    device_class: str = cv.UNDEFINED,
+    entity_category: str = cv.UNDEFINED,
+) -> cv.Schema:
+    schema = {}
+
+    if class_ is not cv.UNDEFINED:
+        schema[cv.GenerateID()] = cv.declare_id(class_)
+
+    for key, default, validator in [
+        (CONF_ICON, icon, cv.icon),
+        (CONF_DEVICE_CLASS, device_class, cv.one_of(*DEVICE_CLASSES, lower=True)),
+        (CONF_ENTITY_CATEGORY, entity_category, cv.entity_category),
+    ]:
+        if default is not cv.UNDEFINED:
+            schema[cv.Optional(key, default=default)] = validator
+
+    return _UPDATE_SCHEMA.extend(schema)
+
+
 async def setup_update_core_(var, config):
     await setup_entity(var, config)