From dc2fe66f2997a468174aef15f07eb5712197f9f1 Mon Sep 17 00:00:00 2001
From: gadgetmobile <57815233+gadgetmobile@users.noreply.github.com>
Date: Tue, 26 May 2020 13:29:19 +0200
Subject: [PATCH] Clean up blebox climate (#36143)

---
 homeassistant/components/blebox/climate.py | 9 +++++----
 tests/components/blebox/test_climate.py    | 3 +++
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/homeassistant/components/blebox/climate.py b/homeassistant/components/blebox/climate.py
index 3f1b21ff687..4ee8cf9be76 100644
--- a/homeassistant/components/blebox/climate.py
+++ b/homeassistant/components/blebox/climate.py
@@ -1,6 +1,6 @@
 """BleBox climate entity."""
 
-from homeassistant.components.climate import ClimateDevice
+from homeassistant.components.climate import ClimateEntity
 from homeassistant.components.climate.const import (
     CURRENT_HVAC_HEAT,
     CURRENT_HVAC_IDLE,
@@ -22,7 +22,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
     )
 
 
-class BleBoxClimateEntity(BleBoxEntity, ClimateDevice):
+class BleBoxClimateEntity(BleBoxEntity, ClimateEntity):
     """Representation of a BleBox climate feature (saunaBox)."""
 
     @property
@@ -81,9 +81,10 @@ class BleBoxClimateEntity(BleBoxEntity, ClimateDevice):
     async def async_set_hvac_mode(self, hvac_mode):
         """Set the climate entity mode."""
         if hvac_mode == HVAC_MODE_HEAT:
-            return await self._feature.async_on()
+            await self._feature.async_on()
+            return
 
-        return await self._feature.async_off()
+        await self._feature.async_off()
 
     async def async_set_temperature(self, **kwargs):
         """Set the thermostat temperature."""
diff --git a/tests/components/blebox/test_climate.py b/tests/components/blebox/test_climate.py
index f7b58cf55c8..c36c93a7f98 100644
--- a/tests/components/blebox/test_climate.py
+++ b/tests/components/blebox/test_climate.py
@@ -135,6 +135,7 @@ async def test_on_when_below_desired(saunabox, hass, config):
         {"entity_id": entity_id, ATTR_HVAC_MODE: HVAC_MODE_HEAT},
         blocking=True,
     )
+    feature_mock.async_off.assert_not_called()
     state = hass.states.get(entity_id)
 
     assert state.attributes[ATTR_HVAC_ACTION] == CURRENT_HVAC_HEAT
@@ -169,6 +170,7 @@ async def test_on_when_above_desired(saunabox, hass, config):
         {"entity_id": entity_id, ATTR_HVAC_MODE: HVAC_MODE_HEAT},
         blocking=True,
     )
+    feature_mock.async_off.assert_not_called()
     state = hass.states.get(entity_id)
 
     assert state.attributes[ATTR_TEMPERATURE] == 23.4
@@ -203,6 +205,7 @@ async def test_off(saunabox, hass, config):
         {"entity_id": entity_id, ATTR_HVAC_MODE: HVAC_MODE_OFF},
         blocking=True,
     )
+    feature_mock.async_on.assert_not_called()
     state = hass.states.get(entity_id)
 
     assert state.attributes[ATTR_HVAC_ACTION] == CURRENT_HVAC_OFF