Improve new Sonos snapshot/restore (#21509)

* Fine-tune new Sonos snapshot/restore

* Move into class

homeassistant/components/sonos/media_player.py

@@ -197,9 +197,11 @@ def _setup_platform(hass, config, add_entities, discovery_info):
 
         with hass.data[DATA_SONOS].topology_lock:
             if service.service == SERVICE_SNAPSHOT:
-                snapshot(entities, service.data[ATTR_WITH_GROUP])
+                SonosEntity.snapshot_multi(
+                    entities, service.data[ATTR_WITH_GROUP])
             elif service.service == SERVICE_RESTORE:
-                restore(entities, service.data[ATTR_WITH_GROUP])
+                SonosEntity.restore_multi(
+                    entities, service.data[ATTR_WITH_GROUP])
             elif service.service == SERVICE_JOIN:
                 master = [e for e in hass.data[DATA_SONOS].entities
                           if e.entity_id == service.data[ATTR_MASTER]]
@@ -357,6 +359,7 @@ class SonosEntity(MediaPlayerDevice):
         self._favorites = None
         self._soco_snapshot = None
         self._snapshot_group = None
+        self._restore_pending = False
 
         self._set_basic_information()
 
@@ -724,6 +727,9 @@ class SonosEntity(MediaPlayerDevice):
                     pass
 
             if self.unique_id == coordinator_uid:
+                if self._restore_pending:
+                    self.restore()
+
                 sonos_group = []
                 for uid in (coordinator_uid, *slave_uids):
                     entity = _get_entity_from_soco_uid(self.hass, uid)
@@ -974,6 +980,82 @@ class SonosEntity(MediaPlayerDevice):
         self.soco.unjoin()
         self._coordinator = None
 
+    @soco_error()
+    def snapshot(self, with_group):
+        """Snapshot the state of a player."""
+        from pysonos.snapshot import Snapshot
+
+        self._soco_snapshot = Snapshot(self.soco)
+        self._soco_snapshot.snapshot()
+        if with_group:
+            self._snapshot_group = self._sonos_group.copy()
+        else:
+            self._snapshot_group = None
+
+    @soco_error()
+    def restore(self):
+        """Restore a snapshotted state to a player."""
+        from pysonos.exceptions import SoCoException
+
+        try:
+            # pylint: disable=protected-access
+            self.soco._zgs_cache.clear()
+            self._soco_snapshot.restore()
+        except (TypeError, AttributeError, SoCoException) as ex:
+            # Can happen if restoring a coordinator onto a current slave
+            _LOGGER.warning("Error on restore %s: %s", self.entity_id, ex)
+
+        self._soco_snapshot = None
+        self._snapshot_group = None
+        self._restore_pending = False
+
+    @staticmethod
+    def snapshot_multi(entities, with_group):
+        """Snapshot all the entities and optionally their groups."""
+        # pylint: disable=protected-access
+        # Find all affected players
+        entities = set(entities)
+        if with_group:
+            for entity in list(entities):
+                entities.update(entity._sonos_group)
+
+        for entity in entities:
+            entity.snapshot(with_group)
+
+    @staticmethod
+    def restore_multi(entities, with_group):
+        """Restore snapshots for all the entities."""
+        # pylint: disable=protected-access
+        # Find all affected players
+        entities = set(e for e in entities if e._soco_snapshot)
+        if with_group:
+            for entity in [e for e in entities if e._snapshot_group]:
+                entities.update(entity._snapshot_group)
+
+        # Pause all current coordinators
+        for entity in (e for e in entities if e.is_coordinator):
+            if entity.state == STATE_PLAYING:
+                entity.media_pause()
+
+        # Bring back the original group topology
+        if with_group:
+            for entity in (e for e in entities if e._snapshot_group):
+                if entity._snapshot_group[0] == entity:
+                    entity.join(entity._snapshot_group)
+
+        # Restore slaves
+        for entity in (e for e in entities if not e.is_coordinator):
+            entity.restore()
+
+        # Restore coordinators (or delay if moving from slave)
+        for entity in (e for e in entities if e.is_coordinator):
+            if entity._sonos_group[0] == entity:
+                # Was already coordinator
+                entity.restore()
+            else:
+                # Await coordinator role
+                entity._restore_pending = True
+
     @soco_error()
     @soco_coordinator
     def set_sleep_timer(self, sleep_time):
@@ -1033,62 +1115,3 @@ class SonosEntity(MediaPlayerDevice):
             attributes[ATTR_SPEECH_ENHANCE] = self._speech_enhance
 
         return attributes
-
-
-@soco_error()
-def snapshot(entities, with_group):
-    """Snapshot all the entities and optionally their groups."""
-    # pylint: disable=protected-access
-    from pysonos.snapshot import Snapshot
-
-    # Find all affected players
-    entities = set(entities)
-    if with_group:
-        for entity in list(entities):
-            entities.update(entity._sonos_group)
-
-    for entity in entities:
-        entity._soco_snapshot = Snapshot(entity.soco)
-        entity._soco_snapshot.snapshot()
-        if with_group:
-            entity._snapshot_group = entity._sonos_group.copy()
-        else:
-            entity._snapshot_group = None
-
-
-@soco_error()
-def restore(entities, with_group):
-    """Restore snapshots for all the entities."""
-    # pylint: disable=protected-access
-    from pysonos.exceptions import SoCoException
-
-    # Find all affected players
-    entities = set(e for e in entities if e._soco_snapshot)
-    if with_group:
-        for entity in [e for e in entities if e._snapshot_group]:
-            entities.update(entity._snapshot_group)
-
-    # Pause all current coordinators
-    for entity in (e for e in entities if e.is_coordinator):
-        if entity.state == STATE_PLAYING:
-            entity.media_pause()
-
-    # Bring back the original group topology and clear pysonos cache
-    if with_group:
-        for entity in (e for e in entities if e._snapshot_group):
-            if entity._snapshot_group[0] == entity:
-                entity.join(entity._snapshot_group)
-                entity.soco._zgs_cache.clear()
-
-    # Restore slaves, then coordinators
-    slaves = [e for e in entities if not e.is_coordinator]
-    coordinators = [e for e in entities if e.is_coordinator]
-    for entity in slaves + coordinators:
-        try:
-            entity._soco_snapshot.restore()
-        except (TypeError, AttributeError, SoCoException) as ex:
-            # Can happen if restoring a coordinator onto a current slave
-            _LOGGER.warning("Error on restore %s: %s", entity.entity_id, ex)
-
-        entity._soco_snapshot = None
-        entity._snapshot_group = None

tests/components/sonos/test_media_player.py

@@ -328,7 +328,7 @@ class TestSonosMediaPlayer(unittest.TestCase):
         snapshotMock.return_value = True
         entity.soco.group = mock.MagicMock()
         entity.soco.group.members = [e.soco for e in entities]
-        sonos.snapshot(entities, True)
+        sonos.SonosEntity.snapshot_multi(entities, True)
         assert snapshotMock.call_count == 1
         assert snapshotMock.call_args == mock.call()
 
@@ -350,6 +350,6 @@ class TestSonosMediaPlayer(unittest.TestCase):
         entity._snapshot_group = mock.MagicMock()
         entity._snapshot_group.members = [e.soco for e in entities]
         entity._soco_snapshot = Snapshot(entity.soco)
-        sonos.restore(entities, True)
+        sonos.SonosEntity.restore_multi(entities, True)
         assert restoreMock.call_count == 1
         assert restoreMock.call_args == mock.call()