refactor: Refactor entity register manager (#1597)

src/utils/diagnostics.ts

@@ -2,7 +2,8 @@ import { HomeAssistant } from '@dermotduffy/custom-card-helpers';
 import pkg from '../../package.json';
 import { RawFrigateCardConfig } from '../config/types';
 import { getLanguage } from '../localize/localize';
-import { DeviceList, getAllDevices } from './ha/device-registry';
+import { getAllDevices } from './ha/registry/device';
+import { DeviceList } from './ha/registry/device/types';
 
 type FrigateVersions = Record<string, string>;
 

src/utils/ha/entity-state-translation.ts

@@ -1,6 +1,6 @@
 import { computeDomain, HomeAssistant } from '@dermotduffy/custom-card-helpers';
 import { HassEntity } from 'home-assistant-js-websocket';
-import { Entity } from './entity-registry/types';
+import { Entity } from './registry/entity/types';
 
 /**
  * Get the translation of an entity state. Inspired by:

src/utils/ha/registry/cache.ts

@@ -1,7 +1,10 @@
-import { Entity } from './types.js';
+export class RegistryCache<T> {
+  protected _cache: Map<string, T> = new Map();
+  protected _keyCallback: (_data: T) => string;
 
-export class EntityCache {
-  protected _cache: Map<string, Entity> = new Map();
+  constructor(keyCallback: (_data: T) => string) {
+    this._keyCallback = keyCallback;
+  }
 
   /**
    * Determine if the cache has a given entity_id.
@@ -17,11 +20,7 @@ export class EntityCache {
    * @param func A callback function that returns a boolean.
    * @returns The first matching value.
    */
-  // public getFirstMatch(func: (arg: T) => boolean): T | null {
-  //   return [...this._cache.values()].find(func) ?? null;
-  // }
-
-  public getMatches(func: (arg: Entity) => boolean): Entity[] {
+  public getMatches(func: (arg: T) => boolean): T[] {
     return [...this._cache.values()].filter(func);
   }
 
@@ -30,16 +29,16 @@ export class EntityCache {
    * @param id The entity id.
    * @returns The entity for this id.
    */
-  public get(id: string): Entity | undefined {
-    return this._cache.get(id);
+  public get(id: string): T | null {
+    return this._cache.get(id) ?? null;
   }
 
   /**
    * Add a given entity to the cache.
    * @param input The entity.
    */
-  public set(input: Entity | Entity[]): void {
-    const _set = (entity: Entity) => this._cache.set(entity.entity_id, entity);
+  public add(input: T | T[]): void {
+    const _set = (arg: T) => this._cache.set(this._keyCallback(arg), arg);
 
     if (Array.isArray(input)) {
       input.forEach(_set);

src/utils/ha/registry/device/index.ts

@@ -1,14 +1,6 @@
 import { HomeAssistant } from '@dermotduffy/custom-card-helpers';
-import { z } from 'zod';
-import { homeAssistantWSRequest } from '.';
-
-const deviceSchema = z.object({
-  model: z.string().nullable(),
-  config_entries: z.string().array(),
-  manufacturer: z.string().nullable(),
-});
-const deviceListSchema = deviceSchema.array();
-export type DeviceList = z.infer<typeof deviceListSchema>;
+import { homeAssistantWSRequest } from '../..';
+import { DeviceList, deviceListSchema } from './types';
 
 /**
  * Get a list of all entities from the entity registry. May throw.

src/utils/ha/registry/device/types.ts

@@ -0,0 +1,12 @@
+import { z } from 'zod';
+
+export const deviceSchema = z.object({
+  id: z.string(),
+  model: z.string().nullable(),
+  config_entries: z.string().array(),
+  manufacturer: z.string().nullable(),
+});
+export type Device = z.infer<typeof deviceSchema>;
+
+export const deviceListSchema = deviceSchema.array();
+export type DeviceList = z.infer<typeof deviceListSchema>;

src/utils/ha/registry/entity/index.ts

@@ -1,17 +1,21 @@
 import { HomeAssistant } from '@dermotduffy/custom-card-helpers';
-import { homeAssistantWSRequest } from '..';
-import { EntityCache } from './cache';
+import { homeAssistantWSRequest } from '../..';
 import { Entity, EntityList, entitySchema, entityListSchema } from './types.js';
+import { RegistryCache } from '../cache';
+
+export const createEntityRegistryCache = (): RegistryCache<Entity> => {
+  return new RegistryCache<Entity>((entity) => entity.entity_id);
+};
 
 // This class manages interactions with entities, caching results and fetching
 // as necessary. Some calls require every entity to be fetched, which may be
-// non-trivial in size (after which it is cached forever).
+// non-trivial in size (after which they are cached forever).
 
 export class EntityRegistryManager {
-  protected _cache: EntityCache;
+  protected _cache: RegistryCache<Entity>;
   protected _fetchedEntityList = false;
 
-  constructor(cache: EntityCache) {
+  constructor(cache: RegistryCache<Entity>) {
     this._cache = cache;
   }
 
@@ -21,11 +25,16 @@ export class EntityRegistryManager {
       return cachedEntity;
     }
 
-    const entity = await homeAssistantWSRequest<Entity>(hass, entitySchema, {
-      type: 'config/entity_registry/get',
-      entity_id: entityID,
-    });
-    this._cache.set(entity);
+    let entity: Entity | null = null;
+    try {
+      entity = await homeAssistantWSRequest<Entity>(hass, entitySchema, {
+        type: 'config/entity_registry/get',
+        entity_id: entityID,
+      });
+    } catch {
+      return null;
+    }
+    this._cache.add(entity);
     return entity;
   }
 
@@ -43,15 +52,11 @@ export class EntityRegistryManager {
   ): Promise<Map<string, Entity>> {
     const output: Map<string, Entity> = new Map();
     const _storeEntity = async (entityID: string): Promise<void> => {
-      let entity: Entity | null = null;
-      try {
-        entity = await this.getEntity(hass, entityID);
-      } catch {
+      const entity = await this.getEntity(hass, entityID);
+
+      if (entity) {
         // When asked to fetch multiple entities, ignore missing entities (they
         // will just not feature in the output).
-        return;
-      }
-      if (entity) {
         output.set(entityID, entity);
       }
     };
@@ -66,7 +71,7 @@ export class EntityRegistryManager {
     const entityList = await homeAssistantWSRequest<EntityList>(hass, entityListSchema, {
       type: 'config/entity_registry/list',
     });
-    this._cache.set(entityList);
+    this._cache.add(entityList);
     this._fetchedEntityList = true;
   }
 }