Move to Launch Library 2 (#42723)

* Move to Launch Library 2

* isort

* Apply suggestions from code review

* Add const to .coveragerc

homeassistant/components/launch_library/sensor.py

@@ -1,8 +1,9 @@
 """A sensor platform that give you information about the next space launch."""
 from datetime import timedelta
 import logging
+from typing import Optional
 
-from pylaunches.api import Launches
+from pylaunches import PyLaunches, PyLaunchesException
 import voluptuous as vol
 
 from homeassistant.components.sensor import PLATFORM_SCHEMA
@@ -11,12 +12,17 @@ from homeassistant.helpers.aiohttp_client import async_get_clientsession
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity import Entity
 
+from .const import (
+    ATTR_AGENCY,
+    ATTR_AGENCY_COUNTRY_CODE,
+    ATTR_LAUNCH_TIME,
+    ATTR_STREAM,
+    ATTRIBUTION,
+    DEFAULT_NAME,
+)
+
 _LOGGER = logging.getLogger(__name__)
 
-ATTRIBUTION = "Data provided by Launch Library."
-
-DEFAULT_NAME = "Next launch"
-
 SCAN_INTERVAL = timedelta(hours=1)
 
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
@@ -26,59 +32,58 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
 
 async def async_setup_platform(hass, config, async_add_entities, discovery_info=None):
     """Create the launch sensor."""
-
     name = config[CONF_NAME]
-
     session = async_get_clientsession(hass)
-    launches = Launches(hass.loop, session)
-    sensor = [LaunchLibrarySensor(launches, name)]
-    async_add_entities(sensor, True)
+    launches = PyLaunches(session)
+
+    async_add_entities([LaunchLibrarySensor(launches, name)], True)
 
 
 class LaunchLibrarySensor(Entity):
     """Representation of a launch_library Sensor."""
 
-    def __init__(self, launches, name):
+    def __init__(self, launches: PyLaunches, name: str) -> None:
         """Initialize the sensor."""
         self.launches = launches
-        self._attributes = {}
+        self.next_launch = None
         self._name = name
-        self._state = None
 
-    async def async_update(self):
+    async def async_update(self) -> None:
         """Get the latest data."""
-        await self.launches.get_launches()
-        if self.launches.launches is None:
-            _LOGGER.error("No data received")
-            return
         try:
-            data = self.launches.launches[0]
-            self._state = data["name"]
-            self._attributes["launch_time"] = data["start"]
-            self._attributes["agency"] = data["agency"]
-            agency_country_code = data["agency_country_code"]
-            self._attributes["agency_country_code"] = agency_country_code
-            self._attributes["stream"] = data["stream"]
-            self._attributes[ATTR_ATTRIBUTION] = ATTRIBUTION
-        except (KeyError, IndexError) as error:
-            _LOGGER.debug("Error getting data, %s", error)
+            launches = await self.launches.upcoming_launches()
+        except PyLaunchesException as exception:
+            _LOGGER.error("Error getting data, %s", exception)
+        else:
+            if launches:
+                self.next_launch = launches[0]
 
     @property
-    def name(self):
+    def name(self) -> str:
         """Return the name of the sensor."""
         return self._name
 
     @property
-    def state(self):
+    def state(self) -> Optional[str]:
         """Return the state of the sensor."""
-        return self._state
+        if self.next_launch:
+            return self.next_launch.name
+        return None
 
     @property
-    def icon(self):
+    def icon(self) -> str:
         """Return the icon of the sensor."""
         return "mdi:rocket"
 
     @property
-    def device_state_attributes(self):
+    def device_state_attributes(self) -> Optional[dict]:
         """Return attributes for the sensor."""
-        return self._attributes
+        if self.next_launch:
+            return {
+                ATTR_LAUNCH_TIME: self.next_launch.net,
+                ATTR_AGENCY: self.next_launch.launch_service_provider.name,
+                ATTR_AGENCY_COUNTRY_CODE: self.next_launch.pad.location.country_code,
+                ATTR_STREAM: self.next_launch.webcast_live,
+                ATTR_ATTRIBUTION: ATTRIBUTION,
+            }
+        return None