Fix available entities on remote

custom_components/unfoldedcircle/helpers.py

@@ -316,7 +316,9 @@ async def synchronize_dock_password(
                 break
 
 
-def update_config_entities(hass: HomeAssistant, client_id: str, entity_ids: [str]):
+def update_config_entities(
+    hass: HomeAssistant, client_id: str, entity_ids: list[str]
+) -> list[str]:
     """Update registry entry of available entities configured in the remote if changed"""
     existing_entries = hass.config_entries.async_entries(domain=DOMAIN)
     try:
@@ -344,23 +346,23 @@ def update_config_entities(hass: HomeAssistant, client_id: str, entity_ids: [str
                 options = dict(config_entry.options)
                 options["available_entities"] = available_entities
                 _LOGGER.debug(
-                    "Available entities need to be updated in registry as there is "
-                    "a desync with the remote %s. "
-                    "Remote : %s, HA registry : %s",
+                    "Available entities need to be updated in registry as there is a desync with the remote %s. Remote : %s, HA registry : %s",
                     client_id,
                     config_entry.options.get("available_entities", []),
                     available_entities,
                 )
                 hass.config_entries.async_update_entry(config_entry, options=options)
+            return available_entities
         else:
             _LOGGER.debug(
                 "Unfolded circle get states from client %s : no config entry", client_id
             )
+            return []
     except StopIteration:
         _LOGGER.debug(
             "Unfolded circle get states from client %s : no config entry", client_id
         )
-        pass
+        return []
 
 
 class UnableToExtractMacAddress(Exception):

custom_components/unfoldedcircle/websocket.py

@@ -79,21 +79,19 @@ def ws_get_states(
     entity_ids: list[str] = msg.get("data", {}).get("entity_ids", [])
     client_id: str | None = msg.get("data", {}).get("client_id", None)
     entity_states = []
+    available_entities = []
     # If entity_ids list is empty, send all entities
     if len(entity_ids) == 0:
         entity_states = hass.states.async_all()
     else:
         # Check if the registry needs to be updated (available entities unsync)
         if client_id:
-            update_config_entities(hass, client_id, entity_ids)
+            available_entities = update_config_entities(hass, client_id, entity_ids)
         else:
             _LOGGER.debug(
                 "No client ID in the request from remote, cannot update the available entities in HA"
             )
 
-        # Add to the requested list the stored list of entities
-        available_entities = []
-
         # Add the missing available entities (normally the unsubscribed entities) to the get states command
         for entity_id in available_entities:
             if entity_id not in entity_ids:
@@ -373,7 +371,8 @@ def remove_listener() -> None:
                 cancel_callback()
             except Exception:
                 pass
-            self._subscriptions.remove(subscription)
+            if subscription in self._subscriptions:
+                self._subscriptions.remove(subscription)
 
         # Create the new events subscription
         subscription_id = msg["id"]
@@ -428,10 +427,11 @@ def remove_listener() -> None:
             """Remove the listener."""
             try:
                 _LOGGER.debug("UC removed configuration event for remote %s", client_id)
-                cancel_callback()
+                # cancel_callback()
             except Exception:
                 pass
-            self._configurations.remove(configuration)
+            if configuration in self._configurations:
+                self._configurations.remove(configuration)
 
         # Create the new events subscription
         subscription_id = msg["id"]