From 5751fdbe5679f0378e18c75b55913b9e5eff610d Mon Sep 17 00:00:00 2001
From: Bram Kragten <mail@bramkragten.nl>
Date: Wed, 3 Apr 2024 15:18:40 +0200
Subject: [PATCH] Improve entity integration filter (#20372)

---
 .../config/entities/ha-config-entities.ts     | 27 +++++++++++++++----
 1 file changed, 22 insertions(+), 5 deletions(-)

diff --git a/src/panels/config/entities/ha-config-entities.ts b/src/panels/config/entities/ha-config-entities.ts
index 910448ac3a..1ec59d4d14 100644
--- a/src/panels/config/entities/ha-config-entities.ts
+++ b/src/panels/config/entities/ha-config-entities.ts
@@ -88,6 +88,10 @@ import { configSections } from "../ha-panel-config";
 import "../integrations/ha-integration-overflow-menu";
 import { showAddIntegrationDialog } from "../integrations/show-add-integration-dialog";
 import { showLabelDetailDialog } from "../labels/show-dialog-label-detail";
+import {
+  EntitySources,
+  fetchEntitySourcesWithCache,
+} from "../../../data/entity_sources";
 
 export interface StateEntity
   extends Omit<EntityRegistryEntry, "id" | "unique_id"> {
@@ -141,6 +145,8 @@ export class HaConfigEntities extends SubscribeMixin(LitElement) {
   @state()
   _labels!: LabelRegistryEntry[];
 
+  @state() private _entitySources?: EntitySources;
+
   @query("hass-tabs-subpage-data-table", true)
   private _dataTable!: HaTabsSubpageDataTable;
 
@@ -405,10 +411,12 @@ export class HaConfigEntities extends SubscribeMixin(LitElement) {
           const entryIds = entries
             .filter((entry) => filter.value!.includes(entry.domain))
             .map((entry) => entry.entry_id);
+
           filteredEntities = filteredEntities.filter(
             (entity) =>
-              entity.config_entry_id &&
-              entryIds.includes(entity.config_entry_id)
+              filter.value?.includes(entity.platform) ||
+              (entity.config_entry_id &&
+                entryIds.includes(entity.config_entry_id))
           );
           filter.value!.forEach((domain) => filteredDomains.add(domain));
         } else if (key === "ha-filter-labels" && filter.value?.length) {
@@ -807,6 +815,9 @@ ${
       },
     };
     this._setFiltersFromUrl();
+    fetchEntitySourcesWithCache(this.hass).then((sources) => {
+      this._entitySources = sources;
+    });
   }
 
   private _setFiltersFromUrl() {
@@ -865,14 +876,18 @@ ${
     this._filters = {};
   }
 
-  public willUpdate(changedProps: PropertyValues<this>): void {
+  public willUpdate(changedProps: PropertyValues): void {
     super.willUpdate(changedProps);
     const oldHass = changedProps.get("hass");
     let changed = false;
     if (!this.hass || !this._entities) {
       return;
     }
-    if (changedProps.has("hass") || changedProps.has("_entities")) {
+    if (
+      changedProps.has("hass") ||
+      changedProps.has("_entities") ||
+      changedProps.has("_entitySources")
+    ) {
       const stateEntities: StateEntity[] = [];
       const regEntityIds = new Set(
         this._entities.map((entity) => entity.entity_id)
@@ -883,6 +898,7 @@ ${
         }
         if (
           !oldHass ||
+          changedProps.has("_entitySources") ||
           this.hass.states[entityId] !== oldHass.states[entityId]
         ) {
           changed = true;
@@ -890,7 +906,8 @@ ${
         stateEntities.push({
           name: computeStateName(this.hass.states[entityId]),
           entity_id: entityId,
-          platform: computeDomain(entityId),
+          platform:
+            this._entitySources?.[entityId]?.domain || computeDomain(entityId),
           disabled_by: null,
           hidden_by: null,
           area_id: null,