From 9daa431075798ef2d022eb41e793f53bc53a9234 Mon Sep 17 00:00:00 2001
From: ludeeus <joasoe@gmail.com>
Date: Sat, 11 Feb 2023 20:49:54 +0000
Subject: [PATCH] Simplify base entity

---
 .../integration_blueprint/coordinator.py      |  3 +++
 .../integration_blueprint/entity.py           | 27 ++++++-------------
 2 files changed, 11 insertions(+), 19 deletions(-)

diff --git a/custom_components/integration_blueprint/coordinator.py b/custom_components/integration_blueprint/coordinator.py
index e9e4371..59d73b9 100644
--- a/custom_components/integration_blueprint/coordinator.py
+++ b/custom_components/integration_blueprint/coordinator.py
@@ -4,6 +4,7 @@
 from datetime import timedelta
 
 from homeassistant.core import HomeAssistant
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
 from .api import IntegrationBlueprintApiClient
@@ -13,6 +14,8 @@
 class BlueprintDataUpdateCoordinator(DataUpdateCoordinator):
     """Class to manage fetching data from the API."""
 
+    config_entry: ConfigEntry
+
     def __init__(
         self,
         hass: HomeAssistant,
diff --git a/custom_components/integration_blueprint/entity.py b/custom_components/integration_blueprint/entity.py
index 24f0235..62c4755 100644
--- a/custom_components/integration_blueprint/entity.py
+++ b/custom_components/integration_blueprint/entity.py
@@ -2,32 +2,21 @@
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from .const import ATTRIBUTION, DOMAIN, NAME, VERSION
+from .coordinator import BlueprintDataUpdateCoordinator
 
 
 class IntegrationBlueprintEntity(CoordinatorEntity):
-    def __init__(self, coordinator, config_entry):
-        super().__init__(coordinator)
-        self.config_entry = config_entry
+    """BlueprintEntity class."""
 
-    @property
-    def unique_id(self):
-        """Return a unique ID to use for this entity."""
-        return self.config_entry.entry_id
+    _attr_attribution = ATTRIBUTION
 
-    @property
-    def device_info(self):
-        return {
+    def __init__(self, coordinator: BlueprintDataUpdateCoordinator) -> None:
+        """Initialize."""
+        super().__init__(coordinator)
+        self._attr_unique_id = coordinator.config_entry.entry_id
+        self._attr_device_info = {
             "identifiers": {(DOMAIN, self.unique_id)},
             "name": NAME,
             "model": VERSION,
             "manufacturer": NAME,
         }
-
-    @property
-    def extra_state_attributes(self):
-        """Return the state attributes."""
-        return {
-            "attribution": ATTRIBUTION,
-            "id": str(self.coordinator.data.get("id")),
-            "integration": DOMAIN,
-        }