Add Twitch subscription and additional stats

homeassistant/components/twitch/sensor.py

@@ -6,16 +6,25 @@
 import voluptuous as vol
 
 from homeassistant.components.sensor import PLATFORM_SCHEMA
+from homeassistant.const import ATTR_FRIENDLY_NAME
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity import Entity
 
 _LOGGER = logging.getLogger(__name__)
 
 ATTR_GAME = "game"
 ATTR_TITLE = "title"
+ATTR_SUBSCRIPTION = "subscribed"
+ATTR_SUBSCRIPTION_SINCE = "subscribed_since"
+ATTR_SUBSCRIPTION_GIFTED = "subscription_is_gifted"
+ATTR_FOLLOW = "following"
+ATTR_FOLLOW_SINCE = "following_since"
+ATTR_FOLLOWING = "follower"
+ATTR_VIEWS = "views"
 
 CONF_CHANNELS = "channels"
 CONF_CLIENT_ID = "client_id"
+CONF_OAUTH_TOKEN = "oauth_token"
 
 ICON = "mdi:twitch"
 
@@ -26,6 +35,7 @@
     {
         vol.Required(CONF_CLIENT_ID): cv.string,
         vol.Required(CONF_CHANNELS): vol.All(cv.ensure_list, [cv.string]),
+        vol.Optional(CONF_OAUTH_TOKEN): cv.string,
     }
 )
 
@@ -34,29 +44,35 @@ def setup_platform(hass, config, add_entities, discovery_info=None):
     """Set up the Twitch platform."""
     channels = config[CONF_CHANNELS]
     client_id = config[CONF_CLIENT_ID]
-    client = TwitchClient(client_id=client_id)
+    oauth_token = config.get(CONF_OAUTH_TOKEN)
+    client = TwitchClient(client_id, oauth_token)
 
     try:
         client.ingests.get_server_list()
     except HTTPError:
-        _LOGGER.error("Client ID is not valid")
+        _LOGGER.error("Client ID or OAuth token is not valid")
         return
 
-    users = client.users.translate_usernames_to_ids(channels)
+    channel_ids = client.users.translate_usernames_to_ids(channels)
 
-    add_entities([TwitchSensor(user, client) for user in users], True)
+    add_entities([TwitchSensor(channel_id, client) for channel_id in channel_ids], True)
 
 
 class TwitchSensor(Entity):
     """Representation of an Twitch channel."""
 
-    def __init__(self, user, client):
+    def __init__(self, channel, client):
         """Initialize the sensor."""
         self._client = client
-        self._user = user
-        self._channel = self._user.name
-        self._id = self._user.id
-        self._state = self._preview = self._game = self._title = None
+        self._channel = channel
+        self._oauth_enabled = client._oauth_token is not None
+        self._state = None
+        self._preview = None
+        self._game = None
+        self._title = None
+        self._subscription = None
+        self._follow = None
+        self._statistics = None
 
     @property
     def should_poll(self):
@@ -66,7 +82,7 @@ def should_poll(self):
     @property
     def name(self):
         """Return the name of the sensor."""
-        return self._channel
+        return self._channel.display_name
 
     @property
     def state(self):
@@ -81,28 +97,67 @@ def entity_picture(self):
     @property
     def device_state_attributes(self):
         """Return the state attributes."""
+        attr = {
+            ATTR_FRIENDLY_NAME: self._channel.display_name,
+        }
+        attr.update(self._statistics)
+
+        if self._oauth_enabled:
+            attr.update(self._subscription)
+            attr.update(self._follow)
+
         if self._state == STATE_STREAMING:
-            return {ATTR_GAME: self._game, ATTR_TITLE: self._title}
+            attr.update({ATTR_GAME: self._game, ATTR_TITLE: self._title})
+        return attr
 
     @property
     def unique_id(self):
         """Return unique ID for this sensor."""
-        return self._id
+        return self._channel.id
 
     @property
     def icon(self):
         """Icon to use in the frontend, if any."""
         return ICON
 
-    # pylint: disable=no-member
     def update(self):
         """Update device state."""
-        stream = self._client.streams.get_stream_by_user(self._id)
+
+        channel = self._client.channels.get_by_id(self._channel.id)
+
+        self._statistics = {
+            ATTR_FOLLOWING: channel.followers,
+            ATTR_VIEWS: channel.views,
+        }
+        if self._oauth_enabled:
+            user = self._client.users.get()
+
+            try:
+                sub = self._client.users.check_subscribed_to_channel(
+                    user.id, self._channel.id
+                )
+                self._subscription = {
+                    ATTR_SUBSCRIPTION: True,
+                    ATTR_SUBSCRIPTION_SINCE: sub.created_at,
+                    ATTR_SUBSCRIPTION_GIFTED: sub.is_gift,
+                }
+            except HTTPError:
+                self._subscription = {ATTR_SUBSCRIPTION: False}
+
+            try:
+                follow = self._client.users.check_follows_channel(
+                    user.id, self._channel.id
+                )
+                self._follow = {ATTR_FOLLOW: True, ATTR_FOLLOW_SINCE: follow.created_at}
+            except HTTPError:
+                self._follow = {ATTR_FOLLOW: False}
+
+        stream = self._client.streams.get_stream_by_user(self._channel.id)
         if stream:
-            self._game = stream.get("channel").get("game")
-            self._title = stream.get("channel").get("status")
-            self._preview = stream.get("preview").get("medium")
+            self._game = stream.channel.get("game")
+            self._title = stream.channel.get("status")
+            self._preview = stream.preview.get("medium")
             self._state = STATE_STREAMING
         else:
-            self._preview = self._client.users.get_by_id(self._id).get("logo")
+            self._preview = self._channel.logo
             self._state = STATE_OFFLINE

requirements_test_all.txt

@@ -551,6 +551,9 @@ python-miio==0.4.8
 # homeassistant.components.nest
 python-nest==4.1.0
 
+# homeassistant.components.twitch
+python-twitch-client==0.6.0
+
 # homeassistant.components.velbus
 python-velbus==2.0.36
 

tests/components/twitch/__init__.py

@@ -0,0 +1 @@
+"""Tests for the Twitch component."""

tests/components/twitch/test_twitch.py

@@ -0,0 +1,174 @@
+"""The tests for an update of the Twitch component."""
+from unittest.mock import MagicMock, patch
+
+from requests import HTTPError
+from twitch.resources import Channel, Follow, Stream, Subscription, User
+
+from homeassistant.components import sensor
+from homeassistant.setup import async_setup_component
+
+ENTITY_ID = "sensor.channel123"
+CONFIG = {
+    sensor.DOMAIN: {
+        "platform": "twitch",
+        "client_id": "1234",
+        "channels": ["channel123"],
+    }
+}
+CONFIG_WITH_OAUTH = {
+    sensor.DOMAIN: {
+        "platform": "twitch",
+        "client_id": "1234",
+        "channels": ["channel123"],
+        "oauth_token": "9876",
+    }
+}
+
+USER_ID = User({"id": 123, "display_name": "channel123", "logo": "logo.png"})
+STREAM_OBJECT_ONLINE = Stream(
+    {
+        "channel": {"game": "Good Game", "status": "Title"},
+        "preview": {"medium": "stream-medium.png"},
+    }
+)
+CHANNEL_OBJECT = Channel({"followers": 42, "views": 24})
+OAUTH_USER_ID = User({"id": 987})
+SUB_ACTIVE = Subscription({"created_at": "2020-01-20T21:22:42", "is_gift": False})
+FOLLOW_ACTIVE = Follow({"created_at": "2020-01-20T21:22:42"})
+
+
+async def test_init(hass):
+    """Test initial config."""
+
+    channels = MagicMock()
+    channels.get_by_id.return_value = CHANNEL_OBJECT
+    streams = MagicMock()
+    streams.get_stream_by_user.return_value = None
+
+    twitch_mock = MagicMock()
+    twitch_mock.users.translate_usernames_to_ids.return_value = [USER_ID]
+    twitch_mock.channels = channels
+    twitch_mock.streams = streams
+
+    with patch(
+        "homeassistant.components.twitch.sensor.TwitchClient", return_value=twitch_mock
+    ):
+        assert await async_setup_component(hass, sensor.DOMAIN, CONFIG) is True
+
+    sensor_state = hass.states.get(ENTITY_ID)
+    assert sensor_state.state == "offline"
+    assert sensor_state.name == "channel123"
+    assert sensor_state.attributes["icon"] == "mdi:twitch"
+    assert sensor_state.attributes["friendly_name"] == "channel123"
+    assert sensor_state.attributes["views"] == 24
+    assert sensor_state.attributes["follower"] == 42
+
+
+async def test_offline(hass):
+    """Test offline state."""
+
+    twitch_mock = MagicMock()
+    twitch_mock.users.translate_usernames_to_ids.return_value = [USER_ID]
+    twitch_mock.channels.get_by_id.return_value = CHANNEL_OBJECT
+    twitch_mock.streams.get_stream_by_user.return_value = None
+
+    with patch(
+        "homeassistant.components.twitch.sensor.TwitchClient", return_value=twitch_mock,
+    ):
+        assert await async_setup_component(hass, sensor.DOMAIN, CONFIG) is True
+
+    sensor_state = hass.states.get(ENTITY_ID)
+    assert sensor_state.state == "offline"
+    assert sensor_state.attributes["entity_picture"] == "logo.png"
+
+
+async def test_streaming(hass):
+    """Test streaming state."""
+
+    twitch_mock = MagicMock()
+    twitch_mock.users.translate_usernames_to_ids.return_value = [USER_ID]
+    twitch_mock.channels.get_by_id.return_value = CHANNEL_OBJECT
+    twitch_mock.streams.get_stream_by_user.return_value = STREAM_OBJECT_ONLINE
+
+    with patch(
+        "homeassistant.components.twitch.sensor.TwitchClient", return_value=twitch_mock,
+    ):
+        assert await async_setup_component(hass, sensor.DOMAIN, CONFIG) is True
+
+    sensor_state = hass.states.get(ENTITY_ID)
+    assert sensor_state.state == "streaming"
+    assert sensor_state.attributes["entity_picture"] == "stream-medium.png"
+    assert sensor_state.attributes["game"] == "Good Game"
+    assert sensor_state.attributes["title"] == "Title"
+
+
+async def test_oauth_without_sub_and_follow(hass):
+    """Test state with oauth."""
+
+    twitch_mock = MagicMock()
+    twitch_mock.users.translate_usernames_to_ids.return_value = [USER_ID]
+    twitch_mock.channels.get_by_id.return_value = CHANNEL_OBJECT
+    twitch_mock._oauth_token = True  # A replacement for the token
+    twitch_mock.users.get.return_value = OAUTH_USER_ID
+    twitch_mock.users.check_subscribed_to_channel.side_effect = HTTPError()
+    twitch_mock.users.check_follows_channel.side_effect = HTTPError()
+
+    with patch(
+        "homeassistant.components.twitch.sensor.TwitchClient", return_value=twitch_mock,
+    ):
+        assert (
+            await async_setup_component(hass, sensor.DOMAIN, CONFIG_WITH_OAUTH) is True
+        )
+
+    sensor_state = hass.states.get(ENTITY_ID)
+    assert sensor_state.attributes["subscribed"] is False
+    assert sensor_state.attributes["following"] is False
+
+
+async def test_oauth_with_sub(hass):
+    """Test state with oauth and sub."""
+
+    twitch_mock = MagicMock()
+    twitch_mock.users.translate_usernames_to_ids.return_value = [USER_ID]
+    twitch_mock.channels.get_by_id.return_value = CHANNEL_OBJECT
+    twitch_mock._oauth_token = True  # A replacement for the token
+    twitch_mock.users.get.return_value = OAUTH_USER_ID
+    twitch_mock.users.check_subscribed_to_channel.return_value = SUB_ACTIVE
+    twitch_mock.users.check_follows_channel.side_effect = HTTPError()
+
+    with patch(
+        "homeassistant.components.twitch.sensor.TwitchClient", return_value=twitch_mock,
+    ):
+        assert (
+            await async_setup_component(hass, sensor.DOMAIN, CONFIG_WITH_OAUTH) is True
+        )
+
+    sensor_state = hass.states.get(ENTITY_ID)
+    assert sensor_state.attributes["subscribed"] is True
+    assert sensor_state.attributes["subscribed_since"] == "2020-01-20T21:22:42"
+    assert sensor_state.attributes["subscription_is_gifted"] is False
+    assert sensor_state.attributes["following"] is False
+
+
+async def test_oauth_with_follow(hass):
+    """Test state with oauth and follow."""
+
+    twitch_mock = MagicMock()
+    twitch_mock.users.translate_usernames_to_ids.return_value = [USER_ID]
+    twitch_mock.channels.get_by_id.return_value = CHANNEL_OBJECT
+    twitch_mock._oauth_token = True  # A replacement for the token
+    twitch_mock.users.get.return_value = OAUTH_USER_ID
+    twitch_mock.users.check_subscribed_to_channel.side_effect = HTTPError()
+    twitch_mock.users.check_follows_channel.return_value = FOLLOW_ACTIVE
+
+    with patch(
+        "homeassistant.components.twitch.sensor.TwitchClient", return_value=twitch_mock,
+    ):
+        assert (
+            await async_setup_component(hass, sensor.DOMAIN, CONFIG_WITH_OAUTH) is True
+        )
+
+    sensor_state = hass.states.get(ENTITY_ID)
+    assert sensor_state.attributes["subscribed"] is False
+    assert sensor_state.attributes["following"] is True
+    assert sensor_state.attributes["following_since"] == "2020-01-20T21:22:42"