Add repeat feature to HEOS media player (#136180)

homeassistant/components/heos/media_player.py

@@ -4,7 +4,6 @@ from __future__ import annotations
 
 from collections.abc import Awaitable, Callable, Coroutine
 from functools import reduce, wraps
-import logging
 from operator import ior
 from typing import Any
 
@@ -14,6 +13,7 @@ from pyheos import (
     HeosError,
     HeosPlayer,
     PlayState,
+    RepeatType,
     const as heos_const,
 )
 
@@ -26,6 +26,7 @@ from homeassistant.components.media_player import (
     MediaPlayerEntityFeature,
     MediaPlayerState,
     MediaType,
+    RepeatMode,
     async_process_play_media_url,
 )
 from homeassistant.core import HomeAssistant
@@ -48,7 +49,6 @@ BASE_SUPPORTED_FEATURES = (
     | MediaPlayerEntityFeature.VOLUME_SET
     | MediaPlayerEntityFeature.VOLUME_STEP
     | MediaPlayerEntityFeature.CLEAR_PLAYLIST
-    | MediaPlayerEntityFeature.SHUFFLE_SET
     | MediaPlayerEntityFeature.SELECT_SOURCE
     | MediaPlayerEntityFeature.PLAY_MEDIA
     | MediaPlayerEntityFeature.GROUPING
@@ -78,7 +78,12 @@ HA_HEOS_ENQUEUE_MAP = {
     MediaPlayerEnqueue.PLAY: AddCriteriaType.PLAY_NOW,
 }
 
-_LOGGER = logging.getLogger(__name__)
+HEOS_HA_REPEAT_TYPE_MAP = {
+    RepeatType.OFF: RepeatMode.OFF,
+    RepeatType.ON_ALL: RepeatMode.ALL,
+    RepeatType.ON_ONE: RepeatMode.ONE,
+}
+HA_HEOS_REPEAT_TYPE_MAP = {v: k for k, v in HEOS_HA_REPEAT_TYPE_MAP.items()}
 
 
 async def async_setup_entry(
@@ -293,6 +298,13 @@ class HeosMediaPlayer(MediaPlayerEntity):
         """Select input source."""
         await self._source_manager.play_source(source, self._player)
 
+    @catch_action_error("set repeat")
+    async def async_set_repeat(self, repeat: RepeatMode) -> None:
+        """Set repeat mode."""
+        await self._player.set_play_mode(
+            HA_HEOS_REPEAT_TYPE_MAP[repeat], self._player.shuffle
+        )
+
     @catch_action_error("set shuffle")
     async def async_set_shuffle(self, shuffle: bool) -> None:
         """Enable/disable shuffle mode."""
@@ -305,11 +317,17 @@ class HeosMediaPlayer(MediaPlayerEntity):
 
     async def async_update(self) -> None:
         """Update supported features of the player."""
+        self._attr_repeat = HEOS_HA_REPEAT_TYPE_MAP[self._player.repeat]
         controls = self._player.now_playing_media.supported_controls
         current_support = [CONTROL_TO_SUPPORT[control] for control in controls]
         self._attr_supported_features = reduce(
             ior, current_support, BASE_SUPPORTED_FEATURES
         )
+        if self.support_next_track and self.support_previous_track:
+            self._attr_supported_features |= (
+                MediaPlayerEntityFeature.REPEAT_SET
+                | MediaPlayerEntityFeature.SHUFFLE_SET
+            )
 
     @catch_action_error("unjoin player")
     async def async_unjoin_player(self) -> None:

tests/components/heos/snapshots/test_media_player.ambr

@@ -19,13 +19,14 @@
       'media_station': 'Station Name',
       'media_title': 'Song',
       'media_type': 'Station',
+      'repeat': <RepeatMode.OFF: 'off'>,
       'shuffle': False,
       'source_list': list([
         "Today's Hits Radio",
         'Classical MPR (Classical Music)',
         'HEOS Drive - Line In 1',
       ]),
-      'supported_features': <MediaPlayerEntityFeature: 2817597>,
+      'supported_features': <MediaPlayerEntityFeature: 3079741>,
       'volume_level': 0.25,
     }),
     'entity_id': 'media_player.test_player',

tests/components/heos/test_media_player.py

@@ -11,6 +11,7 @@ from pyheos import (
     MediaItem,
     PlayerUpdateResult,
     PlayState,
+    RepeatType,
     SignalHeosEvent,
     SignalType,
     const,
@@ -30,6 +31,7 @@ from homeassistant.components.media_player import (
     ATTR_MEDIA_ENQUEUE,
     ATTR_MEDIA_POSITION,
     ATTR_MEDIA_POSITION_UPDATED_AT,
+    ATTR_MEDIA_REPEAT,
     ATTR_MEDIA_SHUFFLE,
     ATTR_MEDIA_VOLUME_LEVEL,
     ATTR_MEDIA_VOLUME_MUTED,
@@ -40,6 +42,7 @@ from homeassistant.components.media_player import (
     SERVICE_SELECT_SOURCE,
     SERVICE_UNJOIN,
     MediaType,
+    RepeatMode,
 )
 from homeassistant.const import (
     ATTR_ENTITY_ID,
@@ -48,6 +51,7 @@ from homeassistant.const import (
     SERVICE_MEDIA_PLAY,
     SERVICE_MEDIA_PREVIOUS_TRACK,
     SERVICE_MEDIA_STOP,
+    SERVICE_REPEAT_SET,
     SERVICE_SHUFFLE_SET,
     SERVICE_VOLUME_MUTE,
     SERVICE_VOLUME_SET,
@@ -563,6 +567,46 @@ async def test_shuffle_set_error(
     player.set_play_mode.assert_called_once_with(player.repeat, True)
 
 
+async def test_repeat_set(
+    hass: HomeAssistant, config_entry: MockConfigEntry, controller: Heos
+) -> None:
+    """Test the repeat set service."""
+    config_entry.add_to_hass(hass)
+    await hass.config_entries.async_setup(config_entry.entry_id)
+    player = controller.players[1]
+    await hass.services.async_call(
+        MEDIA_PLAYER_DOMAIN,
+        SERVICE_REPEAT_SET,
+        {ATTR_ENTITY_ID: "media_player.test_player", ATTR_MEDIA_REPEAT: RepeatMode.ONE},
+        blocking=True,
+    )
+    player.set_play_mode.assert_called_once_with(RepeatType.ON_ONE, player.shuffle)
+
+
+async def test_repeat_set_error(
+    hass: HomeAssistant, config_entry: MockConfigEntry, controller: Heos
+) -> None:
+    """Test the repeat set service raises error."""
+    config_entry.add_to_hass(hass)
+    await hass.config_entries.async_setup(config_entry.entry_id)
+    player = controller.players[1]
+    player.set_play_mode.side_effect = CommandFailedError(None, "Failure", 1)
+    with pytest.raises(
+        HomeAssistantError,
+        match=re.escape("Unable to set repeat: Failure (1)"),
+    ):
+        await hass.services.async_call(
+            MEDIA_PLAYER_DOMAIN,
+            SERVICE_REPEAT_SET,
+            {
+                ATTR_ENTITY_ID: "media_player.test_player",
+                ATTR_MEDIA_REPEAT: RepeatMode.ALL,
+            },
+            blocking=True,
+        )
+    player.set_play_mode.assert_called_once_with(RepeatType.ON_ALL, player.shuffle)
+
+
 async def test_volume_set(
     hass: HomeAssistant, config_entry: MockConfigEntry, controller: Heos
 ) -> None: