From fb25902de9b4cef707699f48ed3c6d06d204c893 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Thu, 20 Jun 2024 18:35:35 -0500
Subject: [PATCH] Cleanup unifiprotect subscriptions logic (#120049)

---
 homeassistant/components/unifiprotect/data.py  | 18 +++++++++++-------
 .../components/unifiprotect/entity.py          |  4 +---
 2 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/homeassistant/components/unifiprotect/data.py b/homeassistant/components/unifiprotect/data.py
index 75c850702f3..e3e4cbc7f50 100644
--- a/homeassistant/components/unifiprotect/data.py
+++ b/homeassistant/components/unifiprotect/data.py
@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from collections import defaultdict
 from collections.abc import Callable, Iterable
 from datetime import datetime, timedelta
 from functools import partial
@@ -78,7 +79,9 @@ class ProtectData:
         self._entry = entry
         self._hass = hass
         self._update_interval = update_interval
-        self._subscriptions: dict[str, list[Callable[[ProtectDeviceType], None]]] = {}
+        self._subscriptions: defaultdict[
+            str, set[Callable[[ProtectDeviceType], None]]
+        ] = defaultdict(set)
         self._pending_camera_ids: set[str] = set()
         self._unsub_interval: CALLBACK_TYPE | None = None
         self._unsub_websocket: CALLBACK_TYPE | None = None
@@ -302,7 +305,7 @@ class ProtectData:
         )
 
     @callback
-    def async_subscribe_device_id(
+    def async_subscribe(
         self, mac: str, update_callback: Callable[[ProtectDeviceType], None]
     ) -> CALLBACK_TYPE:
         """Add an callback subscriber."""
@@ -310,11 +313,11 @@ class ProtectData:
             self._unsub_interval = async_track_time_interval(
                 self._hass, self._async_poll, self._update_interval
             )
-        self._subscriptions.setdefault(mac, []).append(update_callback)
-        return partial(self.async_unsubscribe_device_id, mac, update_callback)
+        self._subscriptions[mac].add(update_callback)
+        return partial(self._async_unsubscribe, mac, update_callback)
 
     @callback
-    def async_unsubscribe_device_id(
+    def _async_unsubscribe(
         self, mac: str, update_callback: Callable[[ProtectDeviceType], None]
     ) -> None:
         """Remove a callback subscriber."""
@@ -328,9 +331,10 @@ class ProtectData:
     @callback
     def _async_signal_device_update(self, device: ProtectDeviceType) -> None:
         """Call the callbacks for a device_id."""
-        if not (subscriptions := self._subscriptions.get(device.mac)):
+        mac = device.mac
+        if not (subscriptions := self._subscriptions.get(mac)):
             return
-        _LOGGER.debug("Updating device: %s (%s)", device.name, device.mac)
+        _LOGGER.debug("Updating device: %s (%s)", device.name, mac)
         for update_callback in subscriptions:
             update_callback(device)
 
diff --git a/homeassistant/components/unifiprotect/entity.py b/homeassistant/components/unifiprotect/entity.py
index a4179e023b3..adf0d334e0a 100644
--- a/homeassistant/components/unifiprotect/entity.py
+++ b/homeassistant/components/unifiprotect/entity.py
@@ -261,9 +261,7 @@ class BaseProtectEntity(Entity):
         """When entity is added to hass."""
         await super().async_added_to_hass()
         self.async_on_remove(
-            self.data.async_subscribe_device_id(
-                self.device.mac, self._async_updated_event
-            )
+            self.data.async_subscribe(self.device.mac, self._async_updated_event)
         )