diff --git a/homeassistant/components/sensor/sonarr.py b/homeassistant/components/sensor/sonarr.py index d60a26f7092..143fcee0a61 100644 --- a/homeassistant/components/sensor/sonarr.py +++ b/homeassistant/components/sensor/sonarr.py @@ -12,11 +12,10 @@ import requests import voluptuous as vol import homeassistant.helpers.config_validation as cv -from homeassistant.const import (CONF_API_KEY, CONF_HOST, CONF_PORT) -from homeassistant.const import CONF_MONITORED_CONDITIONS -from homeassistant.const import CONF_SSL -from homeassistant.helpers.entity import Entity from homeassistant.components.sensor import PLATFORM_SCHEMA +from homeassistant.const import ( + CONF_API_KEY, CONF_HOST, CONF_PORT, CONF_MONITORED_CONDITIONS, CONF_SSL) +from homeassistant.helpers.entity import Entity _LOGGER = logging.getLogger(__name__) @@ -54,15 +53,15 @@ ENDPOINTS = { BYTE_SIZES = ['B', 'KB', 'MB', 'GB', 'TB', 'PB', 'EB', 'ZB', 'YB'] PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({ vol.Required(CONF_API_KEY): cv.string, - vol.Required(CONF_MONITORED_CONDITIONS, default=[]): - vol.All(cv.ensure_list, [vol.In(list(SENSOR_TYPES.keys()))]), - vol.Optional(CONF_INCLUDED, default=[]): cv.ensure_list, - vol.Optional(CONF_SSL, default=False): cv.boolean, - vol.Optional(CONF_HOST, default=DEFAULT_HOST): cv.string, - vol.Optional(CONF_PORT, default=DEFAULT_PORT): cv.port, - vol.Optional(CONF_URLBASE, default=DEFAULT_URLBASE): cv.string, vol.Optional(CONF_DAYS, default=DEFAULT_DAYS): cv.string, - vol.Optional(CONF_UNIT, default=DEFAULT_UNIT): vol.In(BYTE_SIZES) + vol.Optional(CONF_HOST, default=DEFAULT_HOST): cv.string, + vol.Optional(CONF_INCLUDED, default=[]): cv.ensure_list, + vol.Optional(CONF_MONITORED_CONDITIONS, default=['upcoming']): + vol.All(cv.ensure_list, [vol.In(list(SENSOR_TYPES))]), + vol.Optional(CONF_PORT, default=DEFAULT_PORT): cv.port, + vol.Optional(CONF_SSL, default=False): cv.boolean, + vol.Optional(CONF_UNIT, default=DEFAULT_UNIT): vol.In(BYTE_SIZES), + vol.Optional(CONF_URLBASE, default=DEFAULT_URLBASE): cv.string, }) @@ -70,13 +69,11 @@ def setup_platform(hass, config, add_devices, discovery_info=None): """Set up the Sonarr platform.""" conditions = config.get(CONF_MONITORED_CONDITIONS) add_devices( - [SonarrSensor(hass, config, sensor) for sensor in conditions] - ) - return True + [SonarrSensor(hass, config, sensor) for sensor in conditions], True) class SonarrSensor(Entity): - """Implemention of the Sonarr sensor.""" + """Implementation of the Sonarr sensor.""" def __init__(self, hass, conf, sensor_type): """Create Sonarr entity.""" @@ -86,13 +83,12 @@ class SonarrSensor(Entity): self.port = conf.get(CONF_PORT) self.urlbase = conf.get(CONF_URLBASE) if self.urlbase: - self.urlbase = "%s/" % self.urlbase.strip('/') + self.urlbase = "{}/".format(self.urlbase.strip('/')) self.apikey = conf.get(CONF_API_KEY) self.included = conf.get(CONF_INCLUDED) self.days = int(conf.get(CONF_DAYS)) self.ssl = 's' if conf.get(CONF_SSL) else '' - - # Object data + self._state = None self.data = [] self._tz = timezone(str(hass.config.time_zone)) self.type = sensor_type @@ -102,69 +98,7 @@ class SonarrSensor(Entity): else: self._unit = SENSOR_TYPES[self.type][1] self._icon = SENSOR_TYPES[self.type][2] - - # Update sensor self._available = False - self.update() - - def update(self): - """Update the data for the sensor.""" - start = get_date(self._tz) - end = get_date(self._tz, self.days) - try: - res = requests.get( - ENDPOINTS[self.type].format( - self.ssl, self.host, self.port, - self.urlbase, self.apikey, start, end), - timeout=5) - except OSError: - _LOGGER.error('Host %s is not available', self.host) - self._available = False - self._state = None - return - - if res.status_code == 200: - if self.type in ['upcoming', 'queue', 'series', 'commands']: - if self.days == 1 and self.type == 'upcoming': - # Sonarr API returns an empty array if start and end dates - # are the same, so we need to filter to just today - self.data = list( - filter( - lambda x: x['airDate'] == str(start), - res.json() - ) - ) - else: - self.data = res.json() - self._state = len(self.data) - elif self.type == 'wanted': - data = res.json() - res = requests.get('{}&pageSize={}'.format( - ENDPOINTS[self.type].format( - self.ssl, self.host, self.port, - self.urlbase, self.apikey), - data['totalRecords']), timeout=5) - self.data = res.json()['records'] - self._state = len(self.data) - elif self.type == 'diskspace': - # If included paths are not provided, use all data - if self.included == []: - self.data = res.json() - else: - # Filter to only show lists that are included - self.data = list( - filter( - lambda x: x['path'] in self.included, - res.json() - ) - ) - self._state = '{:.2f}'.format( - to_unit( - sum([data['freeSpace'] for data in self.data]), - self._unit - ) - ) - self._available = True @property def name(self): @@ -193,9 +127,7 @@ class SonarrSensor(Entity): if self.type == 'upcoming': for show in self.data: attributes[show['series']['title']] = 'S{:02d}E{:02d}'.format( - show['seasonNumber'], - show['episodeNumber'] - ) + show['seasonNumber'], show['episodeNumber']) elif self.type == 'queue': for show in self.data: attributes[show['series']['title'] + ' S{:02d}E{:02d}'.format( @@ -223,9 +155,7 @@ class SonarrSensor(Entity): elif self.type == 'series': for show in self.data: attributes[show['title']] = '{}/{} Episodes'.format( - show['episodeFileCount'], - show['episodeCount'] - ) + show['episodeFileCount'], show['episodeCount']) return attributes @property @@ -233,6 +163,62 @@ class SonarrSensor(Entity): """Return the icon of the sensor.""" return self._icon + def update(self): + """Update the data for the sensor.""" + start = get_date(self._tz) + end = get_date(self._tz, self.days) + try: + res = requests.get(ENDPOINTS[self.type].format( + self.ssl, self.host, self.port, self.urlbase, self.apikey, + start, end), timeout=5) + except OSError: + _LOGGER.error("Host %s is not available", self.host) + self._available = False + self._state = None + return + + if res.status_code == 200: + if self.type in ['upcoming', 'queue', 'series', 'commands']: + if self.days == 1 and self.type == 'upcoming': + # Sonarr API returns an empty array if start and end dates + # are the same, so we need to filter to just today + self.data = list( + filter( + lambda x: x['airDate'] == str(start), + res.json() + ) + ) + else: + self.data = res.json() + self._state = len(self.data) + elif self.type == 'wanted': + data = res.json() + res = requests.get('{}&pageSize={}'.format( + ENDPOINTS[self.type].format( + self.ssl, self.host, self.port, self.urlbase, + self.apikey), data['totalRecords']), timeout=5) + self.data = res.json()['records'] + self._state = len(self.data) + elif self.type == 'diskspace': + # If included paths are not provided, use all data + if self.included == []: + self.data = res.json() + else: + # Filter to only show lists that are included + self.data = list( + filter( + lambda x: x['path'] in self.included, + res.json() + ) + ) + self._state = '{:.2f}'.format( + to_unit( + sum([data['freeSpace'] for data in self.data]), + self._unit + ) + ) + self._available = True + def get_date(zone, offset=0): """Get date based on timezone and offset of days.""" diff --git a/tests/components/sensor/test_sonarr.py b/tests/components/sensor/test_sonarr.py index dbf918812cc..b71b96e1400 100644 --- a/tests/components/sensor/test_sonarr.py +++ b/tests/components/sensor/test_sonarr.py @@ -1,4 +1,4 @@ -"""The tests for the sonarr platform.""" +"""The tests for the Sonarr platform.""" import unittest import time from datetime import datetime @@ -561,7 +561,7 @@ class TestSonarrSetup(unittest.TestCase): # pylint: disable=invalid-name DEVICES = [] - def add_devices(self, devices): + def add_devices(self, devices, update): """Mock add devices.""" for device in devices: self.DEVICES.append(device)