From f05caf451ec148726e65cebaf42bf26e0acf5b10 Mon Sep 17 00:00:00 2001
From: Chris Talkington <chris@talkingtontech.com>
Date: Mon, 7 Feb 2022 02:53:23 -0600
Subject: [PATCH] Small cleanup of sonarr sensor platform (#65962)

---
 homeassistant/components/sonarr/sensor.py | 30 +++++++++++++++++------
 tests/components/sonarr/test_sensor.py    |  6 +++++
 2 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/homeassistant/components/sonarr/sensor.py b/homeassistant/components/sonarr/sensor.py
index 91e1eeb3257..01046ded7c2 100644
--- a/homeassistant/components/sonarr/sensor.py
+++ b/homeassistant/components/sonarr/sensor.py
@@ -1,11 +1,14 @@
 """Support for Sonarr sensors."""
 from __future__ import annotations
 
+from collections.abc import Awaitable, Callable, Coroutine
 from datetime import timedelta
+from functools import wraps
 import logging
-from typing import Any
+from typing import Any, TypeVar
 
 from sonarr import Sonarr, SonarrConnectionError, SonarrError
+from typing_extensions import Concatenate, ParamSpec
 
 from homeassistant.components.sensor import SensorEntity, SensorEntityDescription
 from homeassistant.config_entries import ConfigEntry
@@ -64,6 +67,9 @@ SENSOR_TYPES: tuple[SensorEntityDescription, ...] = (
     ),
 )
 
+_T = TypeVar("_T", bound="SonarrSensor")
+_P = ParamSpec("_P")
+
 
 async def async_setup_entry(
     hass: HomeAssistant,
@@ -82,14 +88,17 @@ async def async_setup_entry(
     async_add_entities(entities, True)
 
 
-def sonarr_exception_handler(func):
+def sonarr_exception_handler(
+    func: Callable[Concatenate[_T, _P], Awaitable[None]]  # type: ignore[misc]
+) -> Callable[Concatenate[_T, _P], Coroutine[Any, Any, None]]:  # type: ignore[misc]
     """Decorate Sonarr calls to handle Sonarr exceptions.
 
     A decorator that wraps the passed in function, catches Sonarr errors,
     and handles the availability of the entity.
     """
 
-    async def handler(self, *args, **kwargs):
+    @wraps(func)
+    async def wrapper(self: _T, *args: _P.args, **kwargs: _P.kwargs) -> None:
         try:
             await func(self, *args, **kwargs)
             self.last_update_success = True
@@ -102,12 +111,17 @@ def sonarr_exception_handler(func):
                 _LOGGER.error("Invalid response from API: %s", error)
                 self.last_update_success = False
 
-    return handler
+    return wrapper
 
 
 class SonarrSensor(SonarrEntity, SensorEntity):
     """Implementation of the Sonarr sensor."""
 
+    data: dict[str, Any]
+    last_update_success: bool
+    upcoming_days: int
+    wanted_max_items: int
+
     def __init__(
         self,
         sonarr: Sonarr,
@@ -119,10 +133,10 @@ class SonarrSensor(SonarrEntity, SensorEntity):
         self.entity_description = description
         self._attr_unique_id = f"{entry_id}_{description.key}"
 
-        self.data: dict[str, Any] = {}
-        self.last_update_success: bool = False
-        self.upcoming_days: int = options[CONF_UPCOMING_DAYS]
-        self.wanted_max_items: int = options[CONF_WANTED_MAX_ITEMS]
+        self.data = {}
+        self.last_update_success = True
+        self.upcoming_days = options[CONF_UPCOMING_DAYS]
+        self.wanted_max_items = options[CONF_WANTED_MAX_ITEMS]
 
         super().__init__(
             sonarr=sonarr,
diff --git a/tests/components/sonarr/test_sensor.py b/tests/components/sonarr/test_sensor.py
index 8acf7d5b2c8..cc15376efb1 100644
--- a/tests/components/sonarr/test_sensor.py
+++ b/tests/components/sonarr/test_sensor.py
@@ -68,30 +68,36 @@ async def test_sensors(
     assert state
     assert state.attributes.get(ATTR_ICON) == "mdi:harddisk"
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == DATA_GIGABYTES
+    assert state.attributes.get("C:\\") == "263.10/465.42GB (56.53%)"
     assert state.state == "263.10"
 
     state = hass.states.get("sensor.sonarr_queue")
     assert state
     assert state.attributes.get(ATTR_ICON) == "mdi:download"
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == "Episodes"
+    assert state.attributes.get("The Andy Griffith Show S01E01") == "100.00%"
     assert state.state == "1"
 
     state = hass.states.get("sensor.sonarr_shows")
     assert state
     assert state.attributes.get(ATTR_ICON) == "mdi:television"
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == "Series"
+    assert state.attributes.get("The Andy Griffith Show") == "0/0 Episodes"
     assert state.state == "1"
 
     state = hass.states.get("sensor.sonarr_upcoming")
     assert state
     assert state.attributes.get(ATTR_ICON) == "mdi:television"
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == "Episodes"
+    assert state.attributes.get("Bob's Burgers") == "S04E11"
     assert state.state == "1"
 
     state = hass.states.get("sensor.sonarr_wanted")
     assert state
     assert state.attributes.get(ATTR_ICON) == "mdi:television"
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == "Episodes"
+    assert state.attributes.get("Bob's Burgers S04E11") == "2014-01-26"
+    assert state.attributes.get("The Andy Griffith Show S01E01") == "1960-10-03"
     assert state.state == "2"