Add tests to Atag integration (#35944)

* add tests

* better error handling in dependency

* dont suppress errors

* add support for multiple devices

* add test for Unauthorized status

* raise error on service call failure

.coveragerc

@@ -63,10 +63,6 @@ omit =
     homeassistant/components/arwn/sensor.py
     homeassistant/components/asterisk_cdr/mailbox.py
     homeassistant/components/asterisk_mbox/*
-    homeassistant/components/atag/__init__.py
-    homeassistant/components/atag/climate.py
-    homeassistant/components/atag/sensor.py
-    homeassistant/components/atag/water_heater.py
     homeassistant/components/aten_pe/*
     homeassistant/components/atome/*
     homeassistant/components/aurora_abb_powerone/sensor.py

homeassistant/components/atag/__init__.py

@@ -31,16 +31,14 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
     session = async_get_clientsession(hass)
 
     coordinator = AtagDataUpdateCoordinator(hass, session, entry)
-    try:
     await coordinator.async_refresh()
-    except AtagException:
-        raise ConfigEntryNotReady
-
     if not coordinator.last_update_success:
         raise ConfigEntryNotReady
 
     hass.data.setdefault(DOMAIN, {})
     hass.data[DOMAIN][entry.entry_id] = coordinator
+    if entry.unique_id is None:
+        hass.config_entries.async_update_entry(entry, unique_id=coordinator.atag.id)
 
     for platform in PLATFORMS:
         hass.async_create_task(
@@ -65,9 +63,8 @@ class AtagDataUpdateCoordinator(DataUpdateCoordinator):
         """Update data via library."""
         with async_timeout.timeout(20):
             try:
-                await self.atag.update()
+                if not await self.atag.update():
-                if not self.atag.report:
+                    raise UpdateFailed("No data received")
-                    raise UpdateFailed("No data")
             except AtagException as error:
                 raise UpdateFailed(error)
         return self.atag.report
@@ -121,11 +118,6 @@ class AtagEntity(Entity):
         """Return the polling requirement of the entity."""
         return False
 
-    @property
-    def unit_of_measurement(self):
-        """Return the unit of measurement of this entity, if any."""
-        return self.coordinator.atag.climate.temp_unit
-
     @property
     def available(self):
         """Return True if entity is available."""

homeassistant/components/atag/climate.py

@@ -12,7 +12,7 @@ from homeassistant.components.climate.const import (
     SUPPORT_PRESET_MODE,
     SUPPORT_TARGET_TEMPERATURE,
 )
-from homeassistant.const import ATTR_TEMPERATURE, TEMP_CELSIUS, TEMP_FAHRENHEIT
+from homeassistant.const import ATTR_TEMPERATURE
 
 from . import CLIMATE, DOMAIN, AtagEntity
 
@@ -66,9 +66,7 @@ class AtagThermostat(AtagEntity, ClimateEntity):
     @property
     def temperature_unit(self):
         """Return the unit of measurement."""
-        if self.coordinator.atag.climate.temp_unit in [TEMP_CELSIUS, TEMP_FAHRENHEIT]:
         return self.coordinator.atag.climate.temp_unit
-        return None
 
     @property
     def current_temperature(self) -> Optional[float]:

homeassistant/components/atag/config_flow.py

@@ -1,9 +1,9 @@
 """Config flow for the Atag component."""
-from pyatag import DEFAULT_PORT, AtagException, AtagOne
+import pyatag  # from pyatag import DEFAULT_PORT, AtagException, AtagOne
 import voluptuous as vol
 
 from homeassistant import config_entries
-from homeassistant.const import CONF_DEVICE, CONF_EMAIL, CONF_HOST, CONF_PORT
+from homeassistant.const import CONF_EMAIL, CONF_HOST, CONF_PORT
 from homeassistant.core import callback
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
@@ -12,7 +12,7 @@ from . import DOMAIN  # pylint: disable=unused-import
 DATA_SCHEMA = {
     vol.Required(CONF_HOST): str,
     vol.Optional(CONF_EMAIL): str,
-    vol.Required(CONF_PORT, default=DEFAULT_PORT): vol.Coerce(int),
+    vol.Required(CONF_PORT, default=pyatag.const.DEFAULT_PORT): vol.Coerce(int),
 }
 
 
@@ -25,21 +25,22 @@ class AtagConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     async def async_step_user(self, user_input=None):
         """Handle a flow initialized by the user."""
 
-        if self._async_current_entries():
-            return self.async_abort(reason="already_configured")
-
         if not user_input:
             return await self._show_form()
         session = async_get_clientsession(self.hass)
         try:
-            atag = AtagOne(session=session, **user_input)
+            atag = pyatag.AtagOne(session=session, **user_input)
             await atag.authorize()
             await atag.update(force=True)
 
-        except AtagException:
+        except pyatag.errors.Unauthorized:
+            return await self._show_form({"base": "unauthorized"})
+        except pyatag.errors.AtagException:
             return await self._show_form({"base": "connection_error"})
 
-        user_input.update({CONF_DEVICE: atag.id})
+        await self.async_set_unique_id(atag.id)
+        self._abort_if_unique_id_configured(updates=user_input)
+
         return self.async_create_entry(title=atag.id, data=user_input)
 
     @callback

homeassistant/components/atag/manifest.json

@@ -3,6 +3,6 @@
   "name": "Atag",
   "config_flow": true,
   "documentation": "https://www.home-assistant.io/integrations/atag/",
-  "requirements": ["pyatag==0.3.1.2"],
+  "requirements": ["pyatag==0.3.3.4"],
   "codeowners": ["@MatsNL"]
 }

homeassistant/components/atag/sensor.py

@@ -5,6 +5,8 @@ from homeassistant.const import (
     PRESSURE_BAR,
     TEMP_CELSIUS,
     TEMP_FAHRENHEIT,
+    TIME_HOURS,
+    UNIT_PERCENTAGE,
 )
 
 from . import DOMAIN, AtagEntity
@@ -65,6 +67,8 @@ class AtagSensor(AtagEntity):
             PRESSURE_BAR,
             TEMP_CELSIUS,
             TEMP_FAHRENHEIT,
+            UNIT_PERCENTAGE,
+            TIME_HOURS,
         ]:
             return self.coordinator.data[self._id].measure
         return None

homeassistant/components/atag/strings.json

@@ -12,10 +12,11 @@
       }
     },
     "error": {
+      "unauthorized": "Pairing denied, check device for auth request",
       "connection_error": "Failed to connect, please try again"
     },
     "abort": {
-      "already_configured": "Only one Atag device can be added to Home Assistant"
+      "already_configured": "This device has already been added to HomeAssistant"
     }
   }
 }

homeassistant/components/atag/water_heater.py

@@ -40,9 +40,8 @@ class AtagWaterHeater(AtagEntity, WaterHeaterEntity):
     @property
     def current_operation(self):
         """Return current operation."""
-        if self.coordinator.atag.dhw.status:
+        operation = self.coordinator.atag.dhw.current_operation
-            return STATE_PERFORMANCE
+        return operation if operation in self.operation_list else STATE_OFF
-        return STATE_OFF
 
     @property
     def operation_list(self):

requirements_all.txt

@@ -1212,7 +1212,7 @@ pyalmond==0.0.2
 pyarlo==0.2.3
 
 # homeassistant.components.atag
-pyatag==0.3.1.2
+pyatag==0.3.3.4
 
 # homeassistant.components.netatmo
 pyatmo==3.3.1

requirements_test_all.txt

@@ -518,7 +518,7 @@ pyalmond==0.0.2
 pyarlo==0.2.3
 
 # homeassistant.components.atag
-pyatag==0.3.1.2
+pyatag==0.3.3.4
 
 # homeassistant.components.netatmo
 pyatmo==3.3.1

tests/components/atag/__init__.py

@@ -1 +1,85 @@
-"""Tests for the Atag component."""
+"""Tests for the Atag integration."""
+
+from homeassistant.components.atag import DOMAIN
+from homeassistant.const import CONF_EMAIL, CONF_HOST, CONF_PORT
+from homeassistant.core import HomeAssistant
+
+from tests.common import MockConfigEntry
+from tests.test_util.aiohttp import AiohttpClientMocker
+
+USER_INPUT = {
+    CONF_HOST: "127.0.0.1",
+    CONF_EMAIL: "atag@domain.com",
+    CONF_PORT: 10000,
+}
+UID = "xxxx-xxxx-xxxx_xx-xx-xxx-xxx"
+PAIR_REPLY = {"pair_reply": {"status": {"device_id": UID}, "acc_status": 2}}
+UPDATE_REPLY = {"update_reply": {"status": {"device_id": UID}, "acc_status": 2}}
+RECEIVE_REPLY = {
+    "retrieve_reply": {
+        "status": {"device_id": UID},
+        "report": {
+            "burning_hours": 1000,
+            "room_temp": 20,
+            "outside_temp": 15,
+            "dhw_water_temp": 30,
+            "ch_water_temp": 40,
+            "ch_water_pres": 1.8,
+            "ch_return_temp": 35,
+            "boiler_status": 0,
+            "tout_avg": 12,
+            "details": {"rel_mod_level": 0},
+        },
+        "control": {
+            "ch_control_mode": 0,
+            "ch_mode": 1,
+            "ch_mode_duration": 0,
+            "ch_mode_temp": 12,
+            "dhw_temp_setp": 40,
+            "dhw_mode": 1,
+            "dhw_mode_temp": 150,
+            "weather_status": 8,
+        },
+        "configuration": {
+            "download_url": "http://firmware.atag-one.com:80/R58",
+            "temp_unit": 0,
+            "dhw_max_set": 65,
+            "dhw_min_set": 40,
+        },
+        "acc_status": 2,
+    }
+}
+
+
+async def init_integration(
+    hass: HomeAssistant,
+    aioclient_mock: AiohttpClientMocker,
+    rgbw: bool = False,
+    skip_setup: bool = False,
+) -> MockConfigEntry:
+    """Set up the Atag integration in Home Assistant."""
+
+    aioclient_mock.get(
+        "http://127.0.0.1:10000/retrieve",
+        json=RECEIVE_REPLY,
+        headers={"Content-Type": "application/json"},
+    )
+    aioclient_mock.post(
+        "http://127.0.0.1:10000/update",
+        json=UPDATE_REPLY,
+        headers={"Content-Type": "application/json"},
+    )
+    aioclient_mock.post(
+        "http://127.0.0.1:10000/pair",
+        json=PAIR_REPLY,
+        headers={"Content-Type": "application/json"},
+    )
+
+    entry = MockConfigEntry(domain=DOMAIN, data=USER_INPUT)
+    entry.add_to_hass(hass)
+
+    if not skip_setup:
+        await hass.config_entries.async_setup(entry.entry_id)
+        await hass.async_block_till_done()
+
+    return entry

tests/components/atag/test_climate.py

@@ -0,0 +1,108 @@
+"""Tests for the Atag climate platform."""
+
+from homeassistant.components.atag import CLIMATE, DOMAIN
+from homeassistant.components.climate import (
+    ATTR_HVAC_ACTION,
+    ATTR_HVAC_MODE,
+    ATTR_PRESET_MODE,
+    HVAC_MODE_HEAT,
+    SERVICE_SET_HVAC_MODE,
+    SERVICE_SET_PRESET_MODE,
+    SERVICE_SET_TEMPERATURE,
+)
+from homeassistant.components.climate.const import CURRENT_HVAC_HEAT, PRESET_AWAY
+from homeassistant.components.homeassistant import (
+    DOMAIN as HA_DOMAIN,
+    SERVICE_UPDATE_ENTITY,
+)
+from homeassistant.const import ATTR_ENTITY_ID, ATTR_TEMPERATURE, STATE_UNKNOWN
+from homeassistant.core import HomeAssistant
+from homeassistant.setup import async_setup_component
+
+from tests.async_mock import PropertyMock, patch
+from tests.components.atag import UID, init_integration
+from tests.test_util.aiohttp import AiohttpClientMocker
+
+CLIMATE_ID = f"{CLIMATE}.{DOMAIN}"
+
+
+async def test_climate(
+    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
+) -> None:
+    """Test the creation and values of Atag climate device."""
+    with patch("pyatag.entities.Climate.status"):
+        entry = await init_integration(hass, aioclient_mock)
+        registry = await hass.helpers.entity_registry.async_get_registry()
+
+        assert registry.async_is_registered(CLIMATE_ID)
+        entry = registry.async_get(CLIMATE_ID)
+        assert entry.unique_id == f"{UID}-{CLIMATE}"
+        assert (
+            hass.states.get(CLIMATE_ID).attributes[ATTR_HVAC_ACTION]
+            == CURRENT_HVAC_HEAT
+        )
+
+
+async def test_setting_climate(
+    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
+) -> None:
+    """Test setting the climate device."""
+    await init_integration(hass, aioclient_mock)
+    with patch("pyatag.entities.Climate.set_temp") as mock_set_temp:
+        await hass.services.async_call(
+            CLIMATE,
+            SERVICE_SET_TEMPERATURE,
+            {ATTR_ENTITY_ID: CLIMATE_ID, ATTR_TEMPERATURE: 15},
+            blocking=True,
+        )
+        await hass.async_block_till_done()
+        mock_set_temp.assert_called_once_with(15)
+
+    with patch("pyatag.entities.Climate.set_preset_mode") as mock_set_preset:
+        await hass.services.async_call(
+            CLIMATE,
+            SERVICE_SET_PRESET_MODE,
+            {ATTR_ENTITY_ID: CLIMATE_ID, ATTR_PRESET_MODE: PRESET_AWAY},
+            blocking=True,
+        )
+        await hass.async_block_till_done()
+        mock_set_preset.assert_called_once_with(PRESET_AWAY)
+
+    with patch("pyatag.entities.Climate.set_hvac_mode") as mock_set_hvac:
+        await hass.services.async_call(
+            CLIMATE,
+            SERVICE_SET_HVAC_MODE,
+            {ATTR_ENTITY_ID: CLIMATE_ID, ATTR_HVAC_MODE: HVAC_MODE_HEAT},
+            blocking=True,
+        )
+        await hass.async_block_till_done()
+        mock_set_hvac.assert_called_once_with(HVAC_MODE_HEAT)
+
+
+async def test_incorrect_modes(
+    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker,
+) -> None:
+    """Test incorrect values are handled correctly."""
+    with patch(
+        "pyatag.entities.Climate.hvac_mode",
+        new_callable=PropertyMock(return_value="bug"),
+    ):
+        await init_integration(hass, aioclient_mock)
+        assert hass.states.get(CLIMATE_ID).state == STATE_UNKNOWN
+
+
+async def test_update_service(
+    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
+) -> None:
+    """Test the updater service is called."""
+    await init_integration(hass, aioclient_mock)
+    await async_setup_component(hass, HA_DOMAIN, {})
+    with patch("pyatag.AtagOne.update") as updater:
+        await hass.services.async_call(
+            HA_DOMAIN,
+            SERVICE_UPDATE_ENTITY,
+            {ATTR_ENTITY_ID: CLIMATE_ID},
+            blocking=True,
+        )
+        await hass.async_block_till_done()
+        updater.assert_called_once()

tests/components/atag/test_config_flow.py

@@ -1,20 +1,19 @@
 """Tests for the Atag config flow."""
-from pyatag import AtagException
+from pyatag import errors
 
 from homeassistant import config_entries, data_entry_flow
 from homeassistant.components.atag import DOMAIN
-from homeassistant.const import CONF_DEVICE, CONF_EMAIL, CONF_HOST, CONF_PORT
+from homeassistant.core import HomeAssistant
 
 from tests.async_mock import PropertyMock, patch
-from tests.common import MockConfigEntry
+from tests.components.atag import (
-
+    PAIR_REPLY,
-FIXTURE_USER_INPUT = {
+    RECEIVE_REPLY,
-    CONF_HOST: "127.0.0.1",
+    UID,
-    CONF_EMAIL: "test@domain.com",
+    USER_INPUT,
-    CONF_PORT: 10000,
+    init_integration,
-}
+)
-FIXTURE_COMPLETE_ENTRY = FIXTURE_USER_INPUT.copy()
+from tests.test_util.aiohttp import AiohttpClientMocker
-FIXTURE_COMPLETE_ENTRY[CONF_DEVICE] = "device_identifier"
 
 
 async def test_show_form(hass):
@@ -27,29 +26,31 @@ async def test_show_form(hass):
     assert result["step_id"] == "user"
 
 
-async def test_one_config_allowed(hass):
+async def test_adding_second_device(
+    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
+) -> None:
     """Test that only one Atag configuration is allowed."""
-    MockConfigEntry(domain="atag", data=FIXTURE_USER_INPUT).add_to_hass(hass)
+    await init_integration(hass, aioclient_mock)
-
     result = await hass.config_entries.flow.async_init(
-        DOMAIN, context={"source": config_entries.SOURCE_USER}
+        DOMAIN, context={"source": config_entries.SOURCE_USER}, data=USER_INPUT
     )
 
     assert result["type"] == data_entry_flow.RESULT_TYPE_ABORT
     assert result["reason"] == "already_configured"
+    with patch(
+        "pyatag.AtagOne.id", new_callable=PropertyMock(return_value="secondary_device"),
+    ):
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN, context={"source": config_entries.SOURCE_USER}, data=USER_INPUT
+        )
+    assert result["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY
 
 
 async def test_connection_error(hass):
     """Test we show user form on Atag connection error."""
-
+    with patch("pyatag.AtagOne.authorize", side_effect=errors.AtagException()):
-    with patch(
-        "homeassistant.components.atag.config_flow.AtagOne.authorize",
-        side_effect=AtagException(),
-    ):
         result = await hass.config_entries.flow.async_init(
-            DOMAIN,
+            DOMAIN, context={"source": config_entries.SOURCE_USER}, data=USER_INPUT,
-            context={"source": config_entries.SOURCE_USER},
-            data=FIXTURE_USER_INPUT,
         )
 
     assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
@@ -57,19 +58,30 @@ async def test_connection_error(hass):
     assert result["errors"] == {"base": "connection_error"}
 
 
-async def test_full_flow_implementation(hass):
+async def test_unauthorized(hass):
-    """Test registering an integration and finishing flow works."""
+    """Test we show correct form when Unauthorized error is raised."""
-    with patch("homeassistant.components.atag.AtagOne.authorize",), patch(
+    with patch("pyatag.AtagOne.authorize", side_effect=errors.Unauthorized()):
-        "homeassistant.components.atag.AtagOne.update",
-    ), patch(
-        "homeassistant.components.atag.AtagOne.id",
-        new_callable=PropertyMock(return_value="device_identifier"),
-    ):
         result = await hass.config_entries.flow.async_init(
-            DOMAIN,
+            DOMAIN, context={"source": config_entries.SOURCE_USER}, data=USER_INPUT,
-            context={"source": config_entries.SOURCE_USER},
+        )
-            data=FIXTURE_USER_INPUT,
+    assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
+    assert result["step_id"] == "user"
+    assert result["errors"] == {"base": "unauthorized"}
+
+
+async def test_full_flow_implementation(
+    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
+) -> None:
+    """Test registering an integration and finishing flow works."""
+    aioclient_mock.get(
+        "http://127.0.0.1:10000/retrieve", json=RECEIVE_REPLY,
+    )
+    aioclient_mock.post(
+        "http://127.0.0.1:10000/pair", json=PAIR_REPLY,
+    )
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}, data=USER_INPUT,
     )
     assert result["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY
-        assert result["title"] == FIXTURE_COMPLETE_ENTRY[CONF_DEVICE]
+    assert result["title"] == UID
-        assert result["data"] == FIXTURE_COMPLETE_ENTRY
+    assert result["result"].unique_id == UID

tests/components/atag/test_init.py

@@ -0,0 +1,39 @@
+"""Tests for the ATAG integration."""
+import aiohttp
+
+from homeassistant.components.atag import DOMAIN
+from homeassistant.config_entries import ENTRY_STATE_SETUP_RETRY
+from homeassistant.core import HomeAssistant
+
+from tests.async_mock import patch
+from tests.components.atag import init_integration
+from tests.test_util.aiohttp import AiohttpClientMocker
+
+
+async def test_config_entry_not_ready(
+    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
+) -> None:
+    """Test configuration entry not ready on library error."""
+    aioclient_mock.get("http://127.0.0.1:10000/retrieve", exc=aiohttp.ClientError)
+    entry = await init_integration(hass, aioclient_mock)
+    assert entry.state == ENTRY_STATE_SETUP_RETRY
+
+
+async def test_config_entry_empty_reply(
+    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
+) -> None:
+    """Test configuration entry not ready when library returns False."""
+    with patch("pyatag.AtagOne.update", return_value=False):
+        entry = await init_integration(hass, aioclient_mock)
+        assert entry.state == ENTRY_STATE_SETUP_RETRY
+
+
+async def test_unload_config_entry(
+    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
+) -> None:
+    """Test the ATAG configuration entry unloading."""
+    entry = await init_integration(hass, aioclient_mock)
+    assert hass.data[DOMAIN]
+    await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()
+    assert not hass.data.get(DOMAIN)

tests/components/atag/test_sensors.py

@@ -0,0 +1,21 @@
+"""Tests for the Atag sensor platform."""
+
+from homeassistant.components.atag.sensor import SENSORS
+from homeassistant.core import HomeAssistant
+
+from tests.components.atag import UID, init_integration
+from tests.test_util.aiohttp import AiohttpClientMocker
+
+
+async def test_sensors(
+    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
+) -> None:
+    """Test the creation of ATAG sensors."""
+    entry = await init_integration(hass, aioclient_mock)
+    registry = await hass.helpers.entity_registry.async_get_registry()
+
+    for item in SENSORS:
+        sensor_id = "_".join(f"sensor.{item}".lower().split())
+        assert registry.async_is_registered(sensor_id)
+        entry = registry.async_get(sensor_id)
+        assert entry.unique_id in [f"{UID}-{v}" for v in SENSORS.values()]

tests/components/atag/test_water_heater.py

@@ -0,0 +1,41 @@
+"""Tests for the Atag water heater platform."""
+
+from homeassistant.components.atag import DOMAIN, WATER_HEATER
+from homeassistant.components.water_heater import SERVICE_SET_TEMPERATURE
+from homeassistant.const import ATTR_ENTITY_ID, ATTR_TEMPERATURE
+from homeassistant.core import HomeAssistant
+
+from tests.async_mock import patch
+from tests.components.atag import UID, init_integration
+from tests.test_util.aiohttp import AiohttpClientMocker
+
+WATER_HEATER_ID = f"{WATER_HEATER}.{DOMAIN}"
+
+
+async def test_water_heater(
+    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
+) -> None:
+    """Test the creation of Atag water heater."""
+    with patch("pyatag.entities.DHW.status"):
+        entry = await init_integration(hass, aioclient_mock)
+        registry = await hass.helpers.entity_registry.async_get_registry()
+
+        assert registry.async_is_registered(WATER_HEATER_ID)
+        entry = registry.async_get(WATER_HEATER_ID)
+        assert entry.unique_id == f"{UID}-{WATER_HEATER}"
+
+
+async def test_setting_target_temperature(
+    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
+) -> None:
+    """Test setting the water heater device."""
+    await init_integration(hass, aioclient_mock)
+    with patch("pyatag.entities.DHW.set_temp") as mock_set_temp:
+        await hass.services.async_call(
+            WATER_HEATER,
+            SERVICE_SET_TEMPERATURE,
+            {ATTR_ENTITY_ID: WATER_HEATER_ID, ATTR_TEMPERATURE: 50},
+            blocking=True,
+        )
+        await hass.async_block_till_done()
+        mock_set_temp.assert_called_once_with(50)