Add multi contracts support for Hydroquebec

homeassistant/components/sensor/hydroquebec.py

@@ -21,13 +21,14 @@
 from homeassistant.util import Throttle
 import homeassistant.helpers.config_validation as cv
 
-REQUIREMENTS = ['pyhydroquebec==0.1.1']
+REQUIREMENTS = ['pyhydroquebec==1.0.0']
 
 _LOGGER = logging.getLogger(__name__)
 
 KILOWATT_HOUR = "kWh"  # type: str
 PRICE = "CAD"  # type: str
 DAYS = "days"  # type: str
+CONF_CONTRACT = "contract"  # type: str
 
 DEFAULT_NAME = "HydroQuebec"
 
@@ -64,6 +65,7 @@
         vol.All(cv.ensure_list, [vol.In(SENSOR_TYPES)]),
     vol.Required(CONF_USERNAME): cv.string,
     vol.Required(CONF_PASSWORD): cv.string,
+    vol.Required(CONF_CONTRACT): cv.string,
     vol.Optional(CONF_NAME, default=DEFAULT_NAME): cv.string,
 })
 
@@ -91,10 +93,12 @@ def setup_platform(hass, config, add_devices, discovery_info=None):
 
     username = config.get(CONF_USERNAME)
     password = config.get(CONF_PASSWORD)
+    contract = config.get(CONF_CONTRACT)
 
     try:
-        hydroquebec_data = HydroquebecData(username, password)
-        hydroquebec_data.update()
+        hydroquebec_data = HydroquebecData(username, password, contract)
+        _LOGGER.info("Contract list: %s",
+                     ", ".join(hydroquebec_data.get_contract_list()))
     except requests.exceptions.HTTPError as error:
         _LOGGER.error("Failt login: %s", error)
         return False
@@ -105,7 +109,7 @@ def setup_platform(hass, config, add_devices, discovery_info=None):
     for variable in config[CONF_MONITORED_VARIABLES]:
         sensors.append(HydroQuebecSensor(hydroquebec_data, variable, name))
 
-    add_devices(sensors)
+    add_devices(sensors, True)
 
 
 class HydroQuebecSensor(Entity):
@@ -122,8 +126,6 @@ def __init__(self, hydroquebec_data, sensor_type, name):
         self.hydroquebec_data = hydroquebec_data
         self._state = None
 
-        self.update()
-
     @property
     def name(self):
         """Return the name of the sensor."""
@@ -153,22 +155,34 @@ def update(self):
 class HydroquebecData(object):
     """Get data from HydroQuebec."""
 
-    def __init__(self, username, password):
+    def __init__(self, username, password, contract=None):
         """Initialize the data object."""
         from pyhydroquebec import HydroQuebecClient
         self.client = HydroQuebecClient(username,
                                         password,
                                         REQUESTS_TIMEOUT)
+        self._contract = contract
         self.data = {}
 
-    @Throttle(MIN_TIME_BETWEEN_UPDATES)
-    def update(self):
-        """Get the latest data from HydroQuebec."""
+    def get_contract_list(self):
+        """Return the contract list."""
+        # Fetch data
+        self._fetch_data()
+        return self.client.get_contracts()
+
+    def _fetch_data(self):
+        """Fetch latest data from HydroQuebec."""
         from pyhydroquebec.client import PyHydroQuebecError
         try:
             self.client.fetch_data()
         except PyHydroQuebecError as exp:
             _LOGGER.error("Error on receive last Hydroquebec data: %s", exp)
             return
+
+    @Throttle(MIN_TIME_BETWEEN_UPDATES)
+    def update(self):
+        """Return the latest collected data from HydroQuebec."""
+        # Fetch data
+        self._fetch_data()
         # Update data
-        self.data = self.client.get_data()
+        self.data = self.client.get_data(self._contract)[self._contract]

requirements_all.txt

@@ -503,7 +503,7 @@ pyhik==0.1.0
 pyhomematic==0.1.22
 
 # homeassistant.components.sensor.hydroquebec
-pyhydroquebec==0.1.1
+pyhydroquebec==1.0.0
 
 # homeassistant.components.device_tracker.icloud
 pyicloud==0.9.1