Optional persistence in generic_thermostat

homeassistant/components/climate/generic_thermostat.py

@@ -13,13 +13,14 @@
 from homeassistant.core import DOMAIN as HA_DOMAIN
 from homeassistant.components.climate import (
     STATE_HEAT, STATE_COOL, STATE_IDLE, ClimateDevice, PLATFORM_SCHEMA,
-    STATE_AUTO)
+    STATE_AUTO, ATTR_OPERATION_MODE)
 from homeassistant.const import (
     ATTR_UNIT_OF_MEASUREMENT, STATE_ON, STATE_OFF, ATTR_TEMPERATURE,
     CONF_NAME, ATTR_ENTITY_ID, SERVICE_TURN_ON, SERVICE_TURN_OFF)
 from homeassistant.helpers import condition
 from homeassistant.helpers.event import (
     async_track_state_change, async_track_time_interval)
+from homeassistant.helpers.restore_state import async_get_last_state
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.restore_state import async_get_last_state
 
@@ -40,7 +41,10 @@
 CONF_COLD_TOLERANCE = 'cold_tolerance'
 CONF_HOT_TOLERANCE = 'hot_tolerance'
 CONF_KEEP_ALIVE = 'keep_alive'
+CONF_PERSISTENCE = 'persistence'
 
+ATTR_NONE = 'none'
+ATTR_BOTH = 'both'
 
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
     vol.Required(CONF_HEATER): cv.entity_id,
@@ -57,6 +61,8 @@
     vol.Optional(CONF_TARGET_TEMP): vol.Coerce(float),
     vol.Optional(CONF_KEEP_ALIVE): vol.All(
         cv.time_period, cv.positive_timedelta),
+    vol.Optional(CONF_PERSISTENCE, default=ATTR_NONE): vol.In(
+        [ATTR_NONE, ATTR_BOTH, CONF_TARGET_TEMP, ATTR_OPERATION_MODE]),
 })
 
 
@@ -74,19 +80,20 @@ def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
     cold_tolerance = config.get(CONF_COLD_TOLERANCE)
     hot_tolerance = config.get(CONF_HOT_TOLERANCE)
     keep_alive = config.get(CONF_KEEP_ALIVE)
+    persistence = config.get(CONF_PERSISTENCE)
 
     async_add_devices([GenericThermostat(
         hass, name, heater_entity_id, sensor_entity_id, min_temp, max_temp,
         target_temp, ac_mode, min_cycle_duration, cold_tolerance,
-        hot_tolerance, keep_alive)])
+        hot_tolerance, keep_alive, persistence)])
 
 
 class GenericThermostat(ClimateDevice):
     """Representation of a Generic Thermostat device."""
 
     def __init__(self, hass, name, heater_entity_id, sensor_entity_id,
                  min_temp, max_temp, target_temp, ac_mode, min_cycle_duration,
-                 cold_tolerance, hot_tolerance, keep_alive):
+                 cold_tolerance, hot_tolerance, keep_alive, persistence):
         """Initialize the thermostat."""
         self.hass = hass
         self._name = name
@@ -103,6 +110,7 @@ def __init__(self, hass, name, heater_entity_id, sensor_entity_id,
         self._min_temp = min_temp
         self._max_temp = max_temp
         self._target_temp = target_temp
+        self._persistence = persistence
         self._unit = hass.config.units.temperature_unit
 
         async_track_state_change(
@@ -118,17 +126,6 @@ def __init__(self, hass, name, heater_entity_id, sensor_entity_id,
         if sensor_state:
             self._async_update_temp(sensor_state)
 
-    @asyncio.coroutine
-    def async_added_to_hass(self):
-        """Run when entity about to be added."""
-        # If we have an old state and no target temp, restore
-        if self._target_temp is None:
-            old_state = yield from async_get_last_state(self.hass,
-                                                        self.entity_id)
-            if old_state is not None:
-                self._target_temp = float(
-                    old_state.attributes[ATTR_TEMPERATURE])
-
     @property
     def should_poll(self):
         """Return the polling state."""
@@ -226,6 +223,19 @@ def _async_sensor_changed(self, entity_id, old_state, new_state):
         self._async_control_heating()
         yield from self.async_update_ha_state()
 
+    @asyncio.coroutine
+    def async_added_to_hass(self):
+        """Handle all entity which are about to be added."""
+        state = yield from async_get_last_state(self.hass, self.entity_id)
+        if not state:
+            return
+        if self._persistence in [ATTR_BOTH, CONF_TARGET_TEMP]:
+            self._target_temp = state.attributes[ATTR_TEMPERATURE]
+        if (self._persistence in [ATTR_BOTH, ATTR_OPERATION_MODE] and
+                state.attributes[ATTR_OPERATION_MODE] == STATE_OFF):
+            self.set_operation_mode(STATE_OFF)
+        return
+
     @callback
     def _async_switch_changed(self, entity_id, old_state, new_state):
         """Handle heater switch state changes."""

tests/components/climate/test_generic_thermostat.py

@@ -88,7 +88,6 @@ def test_setup_with_sensor(self):
 
 class TestGenericThermostatHeaterSwitching(unittest.TestCase):
     """Test the Generic thermostat heater switching.
-
     Different toggle type devices are tested.
     """
 
@@ -335,7 +334,6 @@ def test_invalid_operating_mode(self, log_mock):
 
     def test_operating_mode_auto(self):
         """Test change mode from OFF to AUTO.
-
         Switch turns on when temp below setpoint and mode changes.
         """
         climate.set_operation_mode(self.hass, STATE_OFF)
@@ -901,6 +899,7 @@ def test_restore_state(hass):
     """Ensure states are restored on startup."""
     mock_restore_cache(hass, (
         State('climate.test_thermostat', '0', {ATTR_TEMPERATURE: "20"}),
+        State('climate.test_thermostat', '0', {'operation_mode': STATE_OFF}),
     ))
 
     hass.state = CoreState.starting
@@ -911,7 +910,9 @@ def test_restore_state(hass):
             'name': 'test_thermostat',
             'heater': ENT_SWITCH,
             'target_sensor': ENT_SENSOR,
+            'persistence': 'both',
         }})
 
     state = hass.states.get('climate.test_thermostat')
     assert(state.attributes[ATTR_TEMPERATURE] == 20)
+    assert(state.attributes['operation_mode'] == STATE_OFF)