@digital-alchemy/hass 25.8.21 → 25.10.19-beta.0

package/src/services/reference.service.mts

@@ -1,17 +1,18 @@
-import { DOWN, NONE, sleep, TAnyFunction, TServiceParams, UP } from "@digital-alchemy/core";
-import dayjs, { Dayjs } from "dayjs";
-import { Get } from "type-fest";
+import type { TAnyFunction, TServiceParams } from "@digital-alchemy/core";
+import { DOWN, NONE, sleep, UP } from "@digital-alchemy/core";
+import type { Dayjs } from "dayjs";
+import dayjs from "dayjs";
+import type { Get } from "type-fest";
 
-import {
+import type {
   ALL_SERVICE_DOMAINS,
   ByIdProxy,
-  domain,
   ENTITY_STATE,
   HassReferenceService,
-  perf,
   RemoveCallback,
 } from "../helpers/index.mts";
-import {
+import { domain, perf } from "../helpers/index.mts";
+import type {
   ANY_ENTITY,
   HassUniqueIdMapping,
   PICK_ENTITY,

package/src/services/registry.service.mts

@@ -1,6 +1,6 @@
-import { TServiceParams } from "@digital-alchemy/core";
+import type { TServiceParams } from "@digital-alchemy/core";
 
-import {
+import type {
   ConfigEntry,
   HassConfig,
   HassRegistryService,

package/src/services/websocket-api.service.mts

@@ -1,16 +1,12 @@
-import {
-  InternalError,
-  SECOND,
-  sleep,
-  START,
-  TBlackHole,
-  TServiceParams,
-} from "@digital-alchemy/core";
-import dayjs, { Dayjs } from "dayjs";
+import type { TBlackHole, TServiceParams } from "@digital-alchemy/core";
+import { InternalError, SECOND, sleep, START } from "@digital-alchemy/core";
+import type { Dayjs } from "dayjs";
+import dayjs from "dayjs";
 import EventEmitter from "events";
+import type { EmptyObject } from "type-fest";
 import WS from "ws";
 
-import {
+import type {
   ConnectionState,
   EntityUpdateEvent,
   HassWebsocketAPI,
@@ -34,6 +30,8 @@ type WaitingMap = Map<
   }
 >;
 
+type MessageHandlerMap = Map<string, Array<(message: { type: string }) => TBlackHole>>;
+
 export function WebsocketAPI({
   context,
   event,
@@ -55,6 +53,7 @@ export function WebsocketAPI({
   let MESSAGE_TIMESTAMPS: number[] = [];
   let onSocketReady: () => void;
   const waitingCallback: WaitingMap = new Map();
+  const messageHandlers: MessageHandlerMap = new Map();
   const isOld = (date: Dayjs) =>
     is.undefined(date) || date.diff(dayjs(), "s") >= config.hass.RETRY_INTERVAL;
 
@@ -293,6 +292,20 @@ export function WebsocketAPI({
     );
   }
 
+  // #MARK: registerMessageHandler
+  function registerMessageHandler<T extends { type: string }>(
+    type: string,
+    callback: (message: T) => TBlackHole,
+  ) {
+    const handlers = messageHandlers.get(type) ?? [];
+    logger.trace({ type }, "register socket message handler");
+    if (!messageHandlers.has(type)) {
+      messageHandlers.set(type, []);
+    }
+    handlers.push(callback as (message: { type: string }) => TBlackHole);
+    messageHandlers.set(type, handlers);
+  }
+
   // #MARK: countMessage
   function countMessage(): void | never {
     messageCount++;
@@ -403,54 +416,18 @@ export function WebsocketAPI({
    * ## result
    * Response to an outgoing emit
    */
-  async function onMessage(message: SocketMessageDTO) {
-    const id = Number(message.id);
+  async function onMessage<T extends { type: string }>(message: T) {
     setImmediate(() => hass.diagnostics.websocket?.message_received.publish({ message }));
-    switch (message.type) {
-      case "auth_required": {
-        logger.trace({ name: onMessage }, `sending authentication`);
-        void hass.socket.sendMessage({ access_token: config.hass.TOKEN, type: "auth" }, false);
-        return;
-      }
-
-      case "auth_ok": {
-        // * Flag as valid connection
-        logger.trace({ name: onMessage }, `event subscriptions starting`);
-        await hass.socket.sendMessage({ type: "subscribe_events" }, false);
-        onSocketReady?.();
-        event.emit(SOCKET_CONNECTED);
-        return;
-      }
 
-      case "event": {
-        return await onMessageEvent(id, message);
-      }
-
-      // 👾
-      case "pong": {
-        // nothing in particular needs to be done, just don't log an error (default)
-        return;
-      }
-
-      case "result": {
-        return await onMessageResult(id, message);
-      }
-
-      case "auth_invalid": {
-        hass.socket.setConnectionState("invalid");
-        logger.fatal(
-          { message, name: onMessage },
-          "received auth invalid {connecting} => {invalid}",
-        );
-        // ? If you have a use case for making this exit configurable, open a ticket
-        process.exit();
-        return;
-      }
+    const handlers = messageHandlers.get(message.type);
+    if (is.empty(handlers)) {
+      logger.error(`unknown websocket message type: ${message.type}`);
+      return;
+    }
 
-      default: {
-        // Code error probably?
-        logger.error({ name: onMessage }, `unknown websocket message type: ${message.type}`);
-      }
+    // Execute all registered handlers for this message type
+    for (const handler of handlers) {
+      await handler(message as { type: string });
     }
   }
 
@@ -523,11 +500,10 @@ export function WebsocketAPI({
   }
 
   // #MARK: subscribe
-  async function subscribe<EVENT extends string>({
-    event_type,
-    context,
-    exec,
-  }: SocketSubscribeOptions<EVENT>) {
+  async function subscribe<
+    EVENT extends string,
+    PAYLOAD extends Record<string, unknown> = EmptyObject,
+  >({ event_type, context, exec }: SocketSubscribeOptions<EVENT, PAYLOAD>) {
     await hass.socket.sendMessage({ event_type, type: "subscribe_events" });
     return hass.socket.onEvent({
       context,
@@ -553,6 +529,43 @@ export function WebsocketAPI({
     return internal.removeFn(() => event.removeListener(SOCKET_CONNECTED, wrapped));
   }
 
+  lifecycle.onPreInit(() => {
+    // Register all current message handlers
+    hass.socket.registerMessageHandler("auth_required", async () => {
+      logger.trace({ name: onMessage }, `sending authentication`);
+      void hass.socket.sendMessage({ access_token: config.hass.TOKEN, type: "auth" }, false);
+    });
+
+    hass.socket.registerMessageHandler("auth_ok", async () => {
+      // * Flag as valid connection
+      logger.trace({ name: onMessage }, `event subscriptions starting`);
+      await hass.socket.sendMessage({ type: "subscribe_events" }, false);
+      onSocketReady?.();
+      event.emit(SOCKET_CONNECTED);
+    });
+
+    hass.socket.registerMessageHandler("event", async (message: SocketMessageDTO) => {
+      const id = Number(message.id);
+      return await onMessageEvent(id, message);
+    });
+
+    hass.socket.registerMessageHandler("pong", async () => {
+      // nothing in particular needs to be done, just don't log an error (default)
+    });
+
+    hass.socket.registerMessageHandler("result", async (message: SocketMessageDTO) => {
+      const id = Number(message.id);
+      return await onMessageResult(id, message);
+    });
+
+    hass.socket.registerMessageHandler("auth_invalid", async (message: SocketMessageDTO) => {
+      hass.socket.setConnectionState("invalid");
+      logger.fatal({ message, name: onMessage }, "received auth invalid {connecting} => {invalid}");
+      // ? If you have a use case for making this exit configurable, open a ticket
+      process.exit();
+    });
+  });
+
   // #MARK: return object
   return {
     attachScheduledFunctions,
@@ -565,6 +578,7 @@ export function WebsocketAPI({
     onEvent,
     onMessage,
     pauseMessages: false,
+    registerMessageHandler,
     sendMessage,
     setConnectionState,
     socketEvents,

package/src/services/zone.service.mts

@@ -1,14 +1,8 @@
-import { debounce, TServiceParams } from "@digital-alchemy/core";
+import type { TServiceParams } from "@digital-alchemy/core";
+import { debounce } from "@digital-alchemy/core";
 
-import {
-  EARLY_ON_READY,
-  HassZoneService,
-  ManifestItem,
-  perf,
-  ZONE_REGISTRY_UPDATED,
-  ZoneDetails,
-  ZoneOptions,
-} from "../helpers/index.mts";
+import type { HassZoneService, ManifestItem, ZoneDetails, ZoneOptions } from "../helpers/index.mts";
+import { EARLY_ON_READY, perf, ZONE_REGISTRY_UPDATED } from "../helpers/index.mts";
 
 export function Zone({
   config,

package/src/testing/area.spec.mts

@@ -1,9 +1,10 @@
 import { sleep } from "@digital-alchemy/core";
 import { subscribe } from "diagnostics_channel";
 
-import { AREA_REGISTRY_UPDATED, AreaDetails } from "../helpers/index.mts";
+import type { AreaDetails } from "../helpers/index.mts";
+import { AREA_REGISTRY_UPDATED } from "../helpers/index.mts";
 import { hassTestRunner, INTERNAL_MESSAGE } from "../mock_assistant/index.mts";
-import { TAreaId } from "../user.mts";
+import type { TAreaId } from "../user.mts";
 
 const EXAMPLE_AREA = {
   area_id: "empty_area" as TAreaId,

package/src/testing/backup.spec.mts

@@ -1,6 +1,6 @@
 import { SINGLE } from "@digital-alchemy/core";
 
-import { BackupResponse } from "../helpers/index.mts";
+import type { BackupResponse } from "../helpers/index.mts";
 import { hassTestRunner } from "../mock_assistant/index.mts";
 
 describe("Backup", () => {

package/src/testing/config.spec.mts

@@ -1,5 +1,5 @@
 import { env } from "process";
-import { MockInstance } from "vitest";
+import type { MockInstance } from "vitest";
 
 import { hassTestRunner } from "../mock_assistant/index.mts";
 

package/src/testing/device.spec.mts

@@ -2,7 +2,7 @@ import { subscribe } from "node:diagnostics_channel";
 
 import { sleep } from "@digital-alchemy/core";
 
-import { DeviceDetails } from "../helpers/index.mts";
+import type { DeviceDetails } from "../helpers/index.mts";
 import { hassTestRunner } from "../mock_assistant/index.mts";
 
 describe("Device", () => {

package/src/testing/entity.spec.mts

@@ -3,9 +3,9 @@ import { subscribe } from "node:diagnostics_channel";
 import { sleep } from "@digital-alchemy/core";
 import dayjs from "dayjs";
 
-import { ENTITY_STATE } from "../index.mts";
+import type { ENTITY_STATE } from "../index.mts";
 import { hassTestRunner } from "../mock_assistant/index.mts";
-import { ANY_ENTITY } from "../user.mts";
+import type { ANY_ENTITY } from "../user.mts";
 
 describe("Entity", () => {
   afterEach(async () => {

package/src/testing/fetch-api.spec.mts

@@ -1,7 +1,7 @@
 import dayjs from "dayjs";
-import { MockInstance } from "vitest";
+import type { MockInstance } from "vitest";
 
-import { HassConfig } from "../helpers/index.mts";
+import type { HassConfig } from "../helpers/index.mts";
 import { hassTestRunner } from "../mock_assistant/index.mts";
 
 describe("FetchAPI", () => {

package/src/testing/floor.spec.mts

@@ -2,9 +2,10 @@ import { subscribe } from "node:diagnostics_channel";
 
 import { sleep } from "@digital-alchemy/core";
 
-import { FLOOR_REGISTRY_UPDATED, FloorDetails } from "../helpers/index.mts";
+import type { FloorDetails } from "../helpers/index.mts";
+import { FLOOR_REGISTRY_UPDATED } from "../helpers/index.mts";
 import { hassTestRunner } from "../mock_assistant/index.mts";
-import { TFloorId } from "../user.mts";
+import type { TFloorId } from "../user.mts";
 
 describe("Floor", () => {
   const EXAMPLE_FLOOR = {

package/src/testing/id-by.spec.mts

@@ -1,5 +1,5 @@
 import { hassTestRunner } from "../mock_assistant/index.mts";
-import { PICK_ENTITY } from "../user.mts";
+import type { PICK_ENTITY } from "../user.mts";
 
 afterEach(async () => {
   await hassTestRunner.teardown();

package/src/testing/label.spec.mts

@@ -2,9 +2,10 @@ import { subscribe } from "node:diagnostics_channel";
 
 import { sleep } from "@digital-alchemy/core";
 
-import { LABEL_REGISTRY_UPDATED, LabelDefinition } from "../helpers/index.mts";
+import type { LabelDefinition } from "../helpers/index.mts";
+import { LABEL_REGISTRY_UPDATED } from "../helpers/index.mts";
 import { hassTestRunner } from "../mock_assistant/index.mts";
-import { TLabelId } from "../user.mts";
+import type { TLabelId } from "../user.mts";
 
 describe("Label", () => {
   const EXAMPLE_LABEL = {

package/src/testing/ref-by.spec.mts

@@ -1,8 +1,8 @@
 import dayjs from "dayjs";
 
-import { ENTITY_STATE } from "../helpers/index.mts";
+import type { ENTITY_STATE } from "../helpers/index.mts";
 import { hassTestRunner } from "../mock_assistant/index.mts";
-import { ANY_ENTITY } from "../user.mts";
+import type { ANY_ENTITY } from "../user.mts";
 
 describe("References", () => {
   afterEach(async () => {

package/src/testing/websocket.spec.mts

@@ -56,4 +56,160 @@ describe("Websocket", () => {
       });
     });
   });
+
+  describe("Message Handler Registration", () => {
+    it("should register and execute message handlers", async () => {
+      expect.assertions(4);
+      await hassTestRunner.run(({ lifecycle, hass }) => {
+        let handler1Called = false;
+        let handler2Called = false;
+
+        hass.socket.registerMessageHandler("test_message", async message => {
+          handler1Called = true;
+          expect(message.type).toBe("test_message");
+        });
+
+        hass.socket.registerMessageHandler("test_message", async message => {
+          handler2Called = true;
+          expect(message.type).toBe("test_message");
+        });
+
+        lifecycle.onReady(async () => {
+          await hass.socket.onMessage({
+            data: { test: "data" },
+            id: 1,
+            type: "test_message",
+          });
+
+          expect(handler1Called).toBe(true);
+          expect(handler2Called).toBe(true);
+        });
+      });
+    });
+
+    it("should handle generic message types", async () => {
+      expect.assertions(2);
+      await hassTestRunner.run(({ lifecycle, hass }) => {
+        interface CustomMessage {
+          type: string;
+          customProperty: string;
+          id: number;
+        }
+
+        let handlerCalled = false;
+
+        hass.socket.registerMessageHandler<CustomMessage>("custom_type", async message => {
+          handlerCalled = true;
+          expect(message.customProperty).toBe("test_value");
+        });
+
+        lifecycle.onReady(async () => {
+          await hass.socket.onMessage({
+            customProperty: "test_value",
+            id: 1,
+            type: "custom_type",
+          });
+
+          expect(handlerCalled).toBe(true);
+        });
+      });
+    });
+
+    it("should handle unknown message types gracefully", async () => {
+      expect.assertions(1);
+      await hassTestRunner.run(({ lifecycle, hass }) => {
+        let errorOccurred = false;
+
+        lifecycle.onReady(async () => {
+          try {
+            // Use a message type that definitely won't have handlers and won't trigger auth flow
+            await hass.socket.onMessage({
+              id: 1,
+              type: "completely_unknown_type_12345",
+            });
+          } catch {
+            errorOccurred = true;
+          }
+
+          // The function should complete without throwing an error
+          expect(errorOccurred).toBe(false);
+        });
+      });
+    });
+
+    it("should execute multiple handlers for the same message type", async () => {
+      expect.assertions(4);
+      await hassTestRunner.run(({ lifecycle, hass }) => {
+        const executionOrder: string[] = [];
+
+        hass.socket.registerMessageHandler("multi_handler", async message => {
+          executionOrder.push("first");
+          expect(message.type).toBe("multi_handler");
+        });
+
+        hass.socket.registerMessageHandler("multi_handler", async message => {
+          executionOrder.push("second");
+          expect(message.type).toBe("multi_handler");
+        });
+
+        hass.socket.registerMessageHandler("multi_handler", async message => {
+          executionOrder.push("third");
+          expect(message.type).toBe("multi_handler");
+        });
+
+        lifecycle.onReady(async () => {
+          await hass.socket.onMessage({
+            id: 1,
+            type: "multi_handler",
+          });
+
+          expect(executionOrder).toEqual(["first", "second", "third"]);
+        });
+      });
+    });
+
+    it("should handle async handlers correctly", async () => {
+      expect.assertions(2);
+      await hassTestRunner.run(({ lifecycle, hass }) => {
+        let asyncHandlerCompleted = false;
+
+        hass.socket.registerMessageHandler("async_test", async message => {
+          await new Promise(resolve => setTimeout(resolve, 10));
+          asyncHandlerCompleted = true;
+          expect(message.type).toBe("async_test");
+        });
+
+        lifecycle.onReady(async () => {
+          await hass.socket.onMessage({
+            id: 1,
+            type: "async_test",
+          });
+
+          expect(asyncHandlerCompleted).toBe(true);
+        });
+      });
+    });
+
+    it("should work with existing message types", async () => {
+      expect.assertions(2);
+      await hassTestRunner.run(({ lifecycle, hass }) => {
+        let customAuthHandlerCalled = false;
+
+        // Register an additional handler for auth_required
+        hass.socket.registerMessageHandler("auth_required", async message => {
+          customAuthHandlerCalled = true;
+          expect(message.type).toBe("auth_required");
+        });
+
+        lifecycle.onReady(async () => {
+          await hass.socket.onMessage({
+            id: 1,
+            type: "auth_required",
+          });
+
+          expect(customAuthHandlerCalled).toBe(true);
+        });
+      });
+    });
+  });
 });

package/src/testing/zone.spec.mts

@@ -2,7 +2,8 @@ import { subscribe } from "node:diagnostics_channel";
 
 import { sleep } from "@digital-alchemy/core";
 
-import { ZONE_REGISTRY_UPDATED, ZoneDetails } from "../helpers/index.mts";
+import type { ZoneDetails } from "../helpers/index.mts";
+import { ZONE_REGISTRY_UPDATED } from "../helpers/index.mts";
 import { hassTestRunner } from "../mock_assistant/index.mts";
 
 describe("Zone", () => {