Improve tests for Shelly WallDisplay (#110435)

Co-authored-by: Maciej Bieniek <478555+bieniu@users.noreply.github.com>

homeassistant/components/shelly/switch.py

@@ -5,7 +5,13 @@ from dataclasses import dataclass
 from typing import Any, cast
 
 from aioshelly.block_device import Block
-from aioshelly.const import MODEL_2, MODEL_25, MODEL_GAS, RPC_GENERATIONS
+from aioshelly.const import (
+    MODEL_2,
+    MODEL_25,
+    MODEL_GAS,
+    MODEL_WALL_DISPLAY,
+    RPC_GENERATIONS,
+)
 
 from homeassistant.components.automation import automations_with_entity
 from homeassistant.components.script import scripts_with_entity
@@ -20,7 +26,7 @@ from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.issue_registry import IssueSeverity, async_create_issue
 
-from .const import DOMAIN, GAS_VALVE_OPEN_STATES, MODEL_WALL_DISPLAY
+from .const import DOMAIN, GAS_VALVE_OPEN_STATES
 from .coordinator import ShellyBlockCoordinator, ShellyRpcCoordinator, get_entry_data
 from .entity import (
     BlockEntityDescription,

tests/components/shelly/test_climate.py

@@ -2,7 +2,7 @@
 from copy import deepcopy
 from unittest.mock import AsyncMock, Mock, PropertyMock
 
-from aioshelly.const import MODEL_VALVE
+from aioshelly.const import MODEL_VALVE, MODEL_WALL_DISPLAY
 from aioshelly.exceptions import DeviceConnectionError, InvalidAuthError
 import pytest
 
@@ -21,7 +21,8 @@ from homeassistant.components.climate import (
     HVACAction,
     HVACMode,
 )
-from homeassistant.components.shelly.const import DOMAIN, MODEL_WALL_DISPLAY
+from homeassistant.components.shelly.const import DOMAIN
+from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
 from homeassistant.config_entries import SOURCE_REAUTH, ConfigEntryState
 from homeassistant.const import ATTR_ENTITY_ID, ATTR_TEMPERATURE, STATE_UNAVAILABLE
 from homeassistant.core import HomeAssistant, State
@@ -682,3 +683,30 @@ async def test_rpc_climate_hvac_mode_cool(
     state = hass.states.get(ENTITY_ID)
     assert state.state == HVACMode.COOL
     assert state.attributes[ATTR_HVAC_ACTION] == HVACAction.COOLING
+
+
+async def test_wall_display_thermostat_mode(
+    hass: HomeAssistant,
+    mock_rpc_device: Mock,
+    entity_registry: EntityRegistry,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Test Wall Display in thermostat mode."""
+    climate_entity_id = "climate.test_name"
+    switch_entity_id = "switch.test_switch_0"
+
+    await init_integration(hass, 2, model=MODEL_WALL_DISPLAY)
+
+    # the switch entity should be removed
+    assert hass.states.get(switch_entity_id) is None
+    assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
+
+    # the climate entity should be created
+    state = hass.states.get(climate_entity_id)
+    assert state
+    assert state.state == HVACMode.HEAT
+    assert len(hass.states.async_entity_ids(CLIMATE_DOMAIN)) == 1
+
+    entry = entity_registry.async_get(climate_entity_id)
+    assert entry
+    assert entry.unique_id == "123456789ABC-thermostat:0"

tests/components/shelly/test_switch.py

@@ -317,30 +317,15 @@ async def test_block_device_gas_valve(
     assert state.state == STATE_ON  # valve is open
 
 
-async def test_wall_display_thermostat_mode(
-    hass: HomeAssistant,
-    mock_rpc_device: Mock,
-) -> None:
-    """Test Wall Display in thermostat mode."""
-    await init_integration(hass, 2, model=MODEL_WALL_DISPLAY)
-
-    # the switch entity should not be created, only the climate entity
-    assert hass.states.get("switch.test_name") is None
-    assert hass.states.get("climate.test_name")
-
-
 async def test_wall_display_relay_mode(
     hass: HomeAssistant,
     mock_rpc_device: Mock,
+    entity_registry: EntityRegistry,
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    """Test Wall Display in thermostat mode."""
+    """Test Wall Display in relay mode."""
-    entity_id = register_entity(
+    climate_entity_id = "climate.test_name"
-        hass,
+    switch_entity_id = "switch.test_switch_0"
-        CLIMATE_DOMAIN,
-        "test_name",
-        "thermostat:0",
-    )
 
     new_shelly = deepcopy(mock_rpc_device.shelly)
     new_shelly["relay_in_thermostat"] = False
@@ -349,7 +334,18 @@ async def test_wall_display_relay_mode(
     await init_integration(hass, 2, model=MODEL_WALL_DISPLAY)
 
     # the climate entity should be removed
-    assert hass.states.get(entity_id) is None
+    assert hass.states.get(climate_entity_id) is None
+    assert len(hass.states.async_entity_ids(CLIMATE_DOMAIN)) == 0
+
+    # the switch entity should be created
+    state = hass.states.get(switch_entity_id)
+    assert state
+    assert state.state == STATE_ON
+    assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
+
+    entry = entity_registry.async_get(switch_entity_id)
+    assert entry
+    assert entry.unique_id == "123456789ABC-switch:0"
 
 
 async def test_create_issue_valve_switch(