Remove helpers and align coding style in Shelly tests (#140080)

* Cleanup hass.states method in Shelly tests (part 1)

* remove helper functions and align coding style

* missed

* revert unwanted changes

* apply review comment

* apply review comment

* apply review comment

* apply ATTR where missing

* apply walrus

* add missed walrus

* add walrus to entity_registry.async_get

* minor tweak

* align after merge

tests/components/shelly/__init__.py

@@ -143,20 +143,6 @@ def get_entity(
     )
 
 
-def get_entity_state(hass: HomeAssistant, entity_id: str) -> str:
-    """Return entity state."""
-    entity = hass.states.get(entity_id)
-    assert entity
-    return entity.state
-
-
-def get_entity_attribute(hass: HomeAssistant, entity_id: str, attribute: str) -> str:
-    """Return entity attribute."""
-    entity = hass.states.get(entity_id)
-    assert entity
-    return entity.attributes[attribute]
-
-
 def register_device(
     device_registry: DeviceRegistry, config_entry: ConfigEntry
 ) -> DeviceEntry:

tests/components/shelly/test_binary_sensor.py

@@ -39,15 +39,16 @@ async def test_block_binary_sensor(
     entity_id = f"{BINARY_SENSOR_DOMAIN}.test_name_channel_1_overpowering"
     await init_integration(hass, 1)
 
-    assert hass.states.get(entity_id).state == STATE_OFF
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_OFF
 
     monkeypatch.setattr(mock_block_device.blocks[RELAY_BLOCK_ID], "overpower", 1)
     mock_block_device.mock_update()
 
-    assert hass.states.get(entity_id).state == STATE_ON
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_ON
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-relay_0-overpower"
 
 
@@ -61,19 +62,18 @@ async def test_block_binary_sensor_extra_state_attr(
     entity_id = f"{BINARY_SENSOR_DOMAIN}.test_name_gas"
     await init_integration(hass, 1)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_ON
     assert state.attributes.get("detected") == "mild"
 
     monkeypatch.setattr(mock_block_device.blocks[SENSOR_BLOCK_ID], "gas", "none")
     mock_block_device.mock_update()
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_OFF
     assert state.attributes.get("detected") == "none"
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-sensor_0-gas"
 
 
@@ -89,15 +89,16 @@ async def test_block_rest_binary_sensor(
     monkeypatch.setitem(mock_block_device.status, "cloud", {"connected": False})
     await init_integration(hass, 1)
 
-    assert hass.states.get(entity_id).state == STATE_OFF
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_OFF
 
     monkeypatch.setitem(mock_block_device.status["cloud"], "connected", True)
     await mock_rest_update(hass, freezer)
 
-    assert hass.states.get(entity_id).state == STATE_ON
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_ON
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-cloud"
 
 
@@ -115,20 +116,22 @@ async def test_block_rest_binary_sensor_connected_battery_devices(
     monkeypatch.setitem(mock_block_device.settings["coiot"], "update_period", 3600)
     await init_integration(hass, 1, model=MODEL_MOTION)
 
-    assert hass.states.get(entity_id).state == STATE_OFF
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_OFF
 
     monkeypatch.setitem(mock_block_device.status["cloud"], "connected", True)
 
     # Verify no update on fast intervals
     await mock_rest_update(hass, freezer)
-    assert hass.states.get(entity_id).state == STATE_OFF
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_OFF
 
     # Verify update on slow intervals
     await mock_rest_update(hass, freezer, seconds=UPDATE_PERIOD_MULTIPLIER * 3600)
-    assert hass.states.get(entity_id).state == STATE_ON
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_ON
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-cloud"
 
 
@@ -149,15 +152,16 @@ async def test_block_sleeping_binary_sensor(
     mock_block_device.mock_online()
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    assert hass.states.get(entity_id).state == STATE_OFF
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_OFF
 
     monkeypatch.setattr(mock_block_device.blocks[SENSOR_BLOCK_ID], "motion", 1)
     mock_block_device.mock_update()
 
-    assert hass.states.get(entity_id).state == STATE_ON
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_ON
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-sensor_0-motion"
 
 
@@ -183,14 +187,16 @@ async def test_block_restored_sleeping_binary_sensor(
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
-    assert hass.states.get(entity_id).state == STATE_ON
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_ON
 
     # Make device online
     monkeypatch.setattr(mock_block_device, "initialized", True)
     mock_block_device.mock_online()
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    assert hass.states.get(entity_id).state == STATE_OFF
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_OFF
 
 
 async def test_block_restored_sleeping_binary_sensor_no_last_state(
@@ -214,14 +220,16 @@ async def test_block_restored_sleeping_binary_sensor_no_last_state(
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
-    assert hass.states.get(entity_id).state == STATE_UNKNOWN
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_UNKNOWN
 
     # Make device online
     monkeypatch.setattr(mock_block_device, "initialized", True)
     mock_block_device.mock_online()
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    assert hass.states.get(entity_id).state == STATE_OFF
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_OFF
 
 
 async def test_rpc_binary_sensor(
@@ -234,17 +242,18 @@ async def test_rpc_binary_sensor(
     entity_id = f"{BINARY_SENSOR_DOMAIN}.test_cover_0_overpowering"
     await init_integration(hass, 2)
 
-    assert hass.states.get(entity_id).state == STATE_OFF
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_OFF
 
     mutate_rpc_device_status(
         monkeypatch, mock_rpc_device, "cover:0", "errors", "overpower"
     )
     mock_rpc_device.mock_update()
 
-    assert hass.states.get(entity_id).state == STATE_ON
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_ON
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-cover:0-overpower"
 
 
@@ -290,20 +299,22 @@ async def test_rpc_sleeping_binary_sensor(
     mock_rpc_device.mock_online()
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    assert hass.states.get(entity_id).state == STATE_OFF
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_OFF
 
     mutate_rpc_device_status(monkeypatch, mock_rpc_device, "cloud", "connected", True)
     mock_rpc_device.mock_update()
 
-    assert hass.states.get(entity_id).state == STATE_ON
+    assert (state := hass.states.get(entity_id))
-
-    # test external power sensor
-    state = hass.states.get("binary_sensor.test_name_external_power")
-    assert state
     assert state.state == STATE_ON
 
-    entry = entity_registry.async_get("binary_sensor.test_name_external_power")
+    # test external power sensor
-    assert entry
+    assert (state := hass.states.get("binary_sensor.test_name_external_power"))
+    assert state.state == STATE_ON
+
+    assert (
+        entry := entity_registry.async_get("binary_sensor.test_name_external_power")
+    )
     assert entry.unique_id == "123456789ABC-devicepower:0-external_power"
 
 
@@ -331,14 +342,16 @@ async def test_rpc_restored_sleeping_binary_sensor(
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
-    assert hass.states.get(entity_id).state == STATE_ON
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_ON
 
     # Make device online
     monkeypatch.setattr(mock_rpc_device, "initialized", True)
     mock_rpc_device.mock_update()
     await hass.async_block_till_done()
 
-    assert hass.states.get(entity_id).state == STATE_OFF
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_OFF
 
 
 async def test_rpc_restored_sleeping_binary_sensor_no_last_state(
@@ -364,7 +377,8 @@ async def test_rpc_restored_sleeping_binary_sensor_no_last_state(
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
-    assert hass.states.get(entity_id).state == STATE_UNKNOWN
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_UNKNOWN
 
     # Make device online
     monkeypatch.setattr(mock_rpc_device, "initialized", True)
@@ -375,7 +389,8 @@ async def test_rpc_restored_sleeping_binary_sensor_no_last_state(
     mock_rpc_device.mock_update()
     await hass.async_block_till_done()
 
-    assert hass.states.get(entity_id).state == STATE_OFF
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_OFF
 
 
 @pytest.mark.parametrize(
@@ -407,17 +422,17 @@ async def test_rpc_device_virtual_binary_sensor(
 
     await init_integration(hass, 3)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert state
     assert state.state == STATE_ON
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-boolean:203-boolean"
 
     monkeypatch.setitem(mock_rpc_device.status["boolean:203"], "value", False)
     mock_rpc_device.mock_update()
-    assert hass.states.get(entity_id).state == STATE_OFF
+
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_OFF
 
 
 async def test_rpc_remove_virtual_binary_sensor_when_mode_toggle(
@@ -450,8 +465,7 @@ async def test_rpc_remove_virtual_binary_sensor_when_mode_toggle(
     await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
 
-    entry = entity_registry.async_get(entity_id)
+    assert entity_registry.async_get(entity_id) is None
-    assert not entry
 
 
 async def test_rpc_remove_virtual_binary_sensor_when_orphaned(
@@ -475,8 +489,7 @@ async def test_rpc_remove_virtual_binary_sensor_when_orphaned(
     await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
 
-    entry = entity_registry.async_get(entity_id)
+    assert entity_registry.async_get(entity_id) is None
-    assert not entry
 
 
 async def test_blu_trv_binary_sensor_entity(

tests/components/shelly/test_button.py

@@ -27,10 +27,10 @@ async def test_block_button(
     entity_id = "button.test_name_reboot"
 
     # reboot button
-    assert hass.states.get(entity_id).state == STATE_UNKNOWN
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_UNKNOWN
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC_reboot"
 
     await hass.services.async_call(
@@ -54,10 +54,10 @@ async def test_rpc_button(
     entity_id = "button.test_name_reboot"
 
     # reboot button
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state == snapshot(name=f"{entity_id}-state")
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
     assert entry == snapshot(name=f"{entity_id}-entry")
 
     await hass.services.async_call(

tests/components/shelly/test_climate.py

@@ -44,13 +44,7 @@ from homeassistant.helpers.device_registry import DeviceRegistry
 from homeassistant.helpers.entity_registry import EntityRegistry
 from homeassistant.util.unit_system import US_CUSTOMARY_SYSTEM
 
-from . import (
+from . import MOCK_MAC, init_integration, register_device, register_entity
-    MOCK_MAC,
-    get_entity_attribute,
-    init_integration,
-    register_device,
-    register_entity,
-)
 from .conftest import MOCK_STATUS_COAP
 
 from tests.common import mock_restore_cache, mock_restore_cache_with_extra_data
@@ -86,11 +80,9 @@ async def test_climate_hvac_mode(
     await hass.async_block_till_done(wait_background_tasks=True)
 
     # Test initial hvac mode - off
-    state = hass.states.get(ENTITY_ID)
+    assert hass.states.get(ENTITY_ID) == snapshot(name=f"{ENTITY_ID}-state")
-    assert state == snapshot(name=f"{ENTITY_ID}-state")
 
-    entry = entity_registry.async_get(ENTITY_ID)
+    assert entity_registry.async_get(ENTITY_ID) == snapshot(name=f"{ENTITY_ID}-entry")
-    assert entry == snapshot(name=f"{ENTITY_ID}-entry")
 
     # Test set hvac mode heat
     await hass.services.async_call(
@@ -105,7 +97,8 @@ async def test_climate_hvac_mode(
 
     monkeypatch.setattr(mock_block_device.blocks[SENSOR_BLOCK_ID], "targetTemp", 20.0)
     mock_block_device.mock_update()
-    state = hass.states.get(ENTITY_ID)
+
+    assert (state := hass.states.get(ENTITY_ID))
     assert state.state == HVACMode.HEAT
 
     # Test set hvac mode off
@@ -122,13 +115,13 @@ async def test_climate_hvac_mode(
 
     monkeypatch.setattr(mock_block_device.blocks[SENSOR_BLOCK_ID], "targetTemp", 4.0)
     mock_block_device.mock_update()
-    state = hass.states.get(ENTITY_ID)
+    assert (state := hass.states.get(ENTITY_ID))
     assert state.state == HVACMode.OFF
 
     # Test unavailable on error
     monkeypatch.setattr(mock_block_device.blocks[DEVICE_BLOCK_ID], "valveError", 1)
     mock_block_device.mock_update()
-    state = hass.states.get(ENTITY_ID)
+    assert (state := hass.states.get(ENTITY_ID))
     assert state.state == STATE_UNAVAILABLE
 
 
@@ -145,7 +138,7 @@ async def test_climate_set_temperature(
     mock_block_device.mock_online()
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    state = hass.states.get(ENTITY_ID)
+    assert (state := hass.states.get(ENTITY_ID))
     assert state.state == HVACMode.OFF
     assert state.attributes[ATTR_TEMPERATURE] == 4
 
@@ -199,7 +192,7 @@ async def test_climate_set_preset_mode(
     mock_block_device.mock_online()
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    state = hass.states.get(ENTITY_ID)
+    assert (state := hass.states.get(ENTITY_ID))
     assert state.attributes[ATTR_PRESET_MODE] == PRESET_NONE
 
     # Test set Profile2
@@ -217,7 +210,7 @@ async def test_climate_set_preset_mode(
     monkeypatch.setattr(mock_block_device.blocks[DEVICE_BLOCK_ID], "mode", 2)
     mock_block_device.mock_update()
 
-    state = hass.states.get(ENTITY_ID)
+    assert (state := hass.states.get(ENTITY_ID))
     assert state.attributes[ATTR_PRESET_MODE] == "Profile2"
 
     # Set preset to none
@@ -236,7 +229,7 @@ async def test_climate_set_preset_mode(
     monkeypatch.setattr(mock_block_device.blocks[DEVICE_BLOCK_ID], "mode", 0)
     mock_block_device.mock_update()
 
-    state = hass.states.get(ENTITY_ID)
+    assert (state := hass.states.get(ENTITY_ID))
     assert state.attributes[ATTR_PRESET_MODE] == PRESET_NONE
 
 
@@ -271,23 +264,26 @@ async def test_block_restored_climate(
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
-    assert hass.states.get(entity_id).state == HVACMode.OFF
+    assert (state := hass.states.get(entity_id))
-    assert hass.states.get(entity_id).attributes.get("temperature") == 4.0
+    assert state.state == HVACMode.OFF
+    assert state.attributes.get(ATTR_TEMPERATURE) == 4.0
 
     # Partial update, should not change state
     mock_block_device.mock_update()
     await hass.async_block_till_done()
 
-    assert hass.states.get(entity_id).state == HVACMode.OFF
+    assert (state := hass.states.get(entity_id))
-    assert hass.states.get(entity_id).attributes.get("temperature") == 4.0
+    assert state.state == HVACMode.OFF
+    assert state.attributes.get(ATTR_TEMPERATURE) == 4.0
 
     # Make device online
     monkeypatch.setattr(mock_block_device, "initialized", True)
     mock_block_device.mock_online()
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    assert hass.states.get(entity_id).state == HVACMode.OFF
+    assert (state := hass.states.get(entity_id))
-    assert hass.states.get(entity_id).attributes.get("temperature") == 4.0
+    assert state.state == HVACMode.OFF
+    assert state.attributes.get(ATTR_TEMPERATURE) == 4.0
 
     # Test set hvac mode heat, target temp should be set to last target temp (22)
     await hass.services.async_call(
@@ -302,9 +298,10 @@ async def test_block_restored_climate(
 
     monkeypatch.setattr(mock_block_device.blocks[SENSOR_BLOCK_ID], "targetTemp", 22.0)
     mock_block_device.mock_update()
-    state = hass.states.get(ENTITY_ID)
+
+    assert (state := hass.states.get(entity_id))
     assert state.state == HVACMode.HEAT
-    assert hass.states.get(entity_id).attributes.get("temperature") == 22.0
+    assert state.attributes.get(ATTR_TEMPERATURE) == 22.0
 
 
 async def test_block_restored_climate_us_customary(
@@ -339,17 +336,19 @@ async def test_block_restored_climate_us_customary(
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
-    assert hass.states.get(entity_id).state == HVACMode.OFF
+    assert (state := hass.states.get(entity_id))
-    assert hass.states.get(entity_id).attributes.get("temperature") == 39
+    assert state.state == HVACMode.OFF
-    assert hass.states.get(entity_id).attributes.get("current_temperature") == 67
+    assert state.attributes.get(ATTR_TEMPERATURE) == 39
+    assert state.attributes.get(ATTR_CURRENT_TEMPERATURE) == 67
 
     # Partial update, should not change state
     mock_block_device.mock_update()
     await hass.async_block_till_done()
 
-    assert hass.states.get(entity_id).state == HVACMode.OFF
+    assert (state := hass.states.get(entity_id))
-    assert hass.states.get(entity_id).attributes.get("temperature") == 39
+    assert state.state == HVACMode.OFF
-    assert hass.states.get(entity_id).attributes.get("current_temperature") == 67
+    assert state.attributes.get(ATTR_TEMPERATURE) == 39
+    assert state.attributes.get(ATTR_CURRENT_TEMPERATURE) == 67
 
     # Make device online
     monkeypatch.setattr(mock_block_device, "initialized", True)
@@ -358,9 +357,10 @@ async def test_block_restored_climate_us_customary(
     mock_block_device.mock_online()
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    assert hass.states.get(entity_id).state == HVACMode.OFF
+    assert (state := hass.states.get(entity_id))
-    assert hass.states.get(entity_id).attributes.get("temperature") == 39
+    assert state.state == HVACMode.OFF
-    assert hass.states.get(entity_id).attributes.get("current_temperature") == 65
+    assert state.attributes.get(ATTR_TEMPERATURE) == 39
+    assert state.attributes.get(ATTR_CURRENT_TEMPERATURE) == 65
 
     # Test set hvac mode heat, target temp should be set to last target temp (10.0/50)
     await hass.services.async_call(
@@ -375,9 +375,10 @@ async def test_block_restored_climate_us_customary(
 
     monkeypatch.setattr(mock_block_device.blocks[SENSOR_BLOCK_ID], "targetTemp", 10.0)
     mock_block_device.mock_update()
-    state = hass.states.get(ENTITY_ID)
+
+    assert (state := hass.states.get(entity_id))
     assert state.state == HVACMode.HEAT
-    assert hass.states.get(entity_id).attributes.get("temperature") == 50
+    assert state.attributes.get(ATTR_TEMPERATURE) == 50
 
 
 async def test_block_restored_climate_unavailable(
@@ -405,7 +406,8 @@ async def test_block_restored_climate_unavailable(
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
-    assert hass.states.get(entity_id).state == HVACMode.OFF
+    assert (state := hass.states.get(entity_id))
+    assert state.state == HVACMode.OFF
 
 
 async def test_block_restored_climate_set_preset_before_online(
@@ -433,7 +435,8 @@ async def test_block_restored_climate_set_preset_before_online(
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
-    assert hass.states.get(entity_id).state == HVACMode.HEAT
+    assert (state := hass.states.get(entity_id))
+    assert state.state == HVACMode.HEAT
 
     with pytest.raises(ServiceValidationError):
         await hass.services.async_call(
@@ -615,16 +618,14 @@ async def test_rpc_climate_hvac_mode(
 
     await init_integration(hass, 2, model=MODEL_WALL_DISPLAY)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id)) == snapshot(name=f"{entity_id}-state")
-    assert state == snapshot(name=f"{entity_id}-state")
 
-    entry = entity_registry.async_get(entity_id)
+    assert entity_registry.async_get(entity_id) == snapshot(name=f"{entity_id}-entry")
-    assert entry == snapshot(name=f"{entity_id}-entry")
 
     monkeypatch.setitem(mock_rpc_device.status["thermostat:0"], "output", False)
     mock_rpc_device.mock_update()
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.attributes[ATTR_HVAC_ACTION] == HVACAction.IDLE
     assert state.attributes[ATTR_CURRENT_HUMIDITY] == 44.4
 
@@ -640,7 +641,7 @@ async def test_rpc_climate_hvac_mode(
     mock_rpc_device.call_rpc.assert_called_once_with(
         "Thermostat.SetConfig", {"config": {"id": 0, "enable": False}}
     )
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == HVACMode.OFF
 
 
@@ -658,15 +659,14 @@ async def test_rpc_climate_without_humidity(
 
     await init_integration(hass, 2, model=MODEL_WALL_DISPLAY)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == HVACMode.HEAT
     assert state.attributes[ATTR_TEMPERATURE] == 23
     assert state.attributes[ATTR_CURRENT_TEMPERATURE] == 12.3
     assert state.attributes[ATTR_HVAC_ACTION] == HVACAction.HEATING
     assert ATTR_CURRENT_HUMIDITY not in state.attributes
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-thermostat:0"
 
 
@@ -678,7 +678,7 @@ async def test_rpc_climate_set_temperature(
 
     await init_integration(hass, 2, model=MODEL_WALL_DISPLAY)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.attributes[ATTR_TEMPERATURE] == 23
 
     monkeypatch.setitem(mock_rpc_device.status["thermostat:0"], "target_C", 28)
@@ -693,7 +693,7 @@ async def test_rpc_climate_set_temperature(
     mock_rpc_device.call_rpc.assert_called_once_with(
         "Thermostat.SetConfig", {"config": {"id": 0, "target_C": 28}}
     )
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.attributes[ATTR_TEMPERATURE] == 28
 
 
@@ -708,7 +708,7 @@ async def test_rpc_climate_hvac_mode_cool(
 
     await init_integration(hass, 2, model=MODEL_WALL_DISPLAY)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == HVACMode.COOL
     assert state.attributes[ATTR_HVAC_ACTION] == HVACAction.COOLING
 
@@ -757,19 +757,16 @@ async def test_wall_display_thermostat_mode_external_actuator(
     await init_integration(hass, 2, model=MODEL_WALL_DISPLAY)
 
     # the switch entity should be created
-    state = hass.states.get(switch_entity_id)
+    assert (state := hass.states.get(switch_entity_id))
-    assert state
     assert state.state == STATE_ON
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
 
     # the climate entity should be created
-    state = hass.states.get(climate_entity_id)
+    assert (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 := entity_registry.async_get(climate_entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-thermostat:0"
 
 
@@ -787,13 +784,9 @@ async def test_blu_trv_climate_set_temperature(
 
     await init_integration(hass, 3, model=MODEL_BLU_GATEWAY_G3)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id)) == snapshot(name=f"{entity_id}-state")
-    assert state == snapshot(name=f"{entity_id}-state")
 
-    entry = entity_registry.async_get(entity_id)
+    assert entity_registry.async_get(entity_id) == snapshot(name=f"{entity_id}-entry")
-    assert entry == snapshot(name=f"{entity_id}-entry")
-
-    assert get_entity_attribute(hass, entity_id, ATTR_TEMPERATURE) == 17.1
 
     monkeypatch.setitem(
         mock_blu_trv.status[f"{BLU_TRV_IDENTIFIER}:200"], "target_C", 28
@@ -816,7 +809,8 @@ async def test_blu_trv_climate_set_temperature(
         BLU_TRV_TIMEOUT,
     )
 
-    assert get_entity_attribute(hass, entity_id, ATTR_TEMPERATURE) == 28
+    assert (state := hass.states.get(entity_id))
+    assert state.attributes[ATTR_TEMPERATURE] == 28
 
 
 async def test_blu_trv_climate_disabled(
@@ -831,14 +825,16 @@ async def test_blu_trv_climate_disabled(
 
     await init_integration(hass, 3, model=MODEL_BLU_GATEWAY_G3)
 
-    assert get_entity_attribute(hass, entity_id, ATTR_TEMPERATURE) == 17.1
+    assert (state := hass.states.get(entity_id))
+    assert state.attributes[ATTR_TEMPERATURE] == 17.1
 
     monkeypatch.setitem(
         mock_blu_trv.config[f"{BLU_TRV_IDENTIFIER}:200"], "enable", False
     )
     mock_blu_trv.mock_update()
 
-    assert get_entity_attribute(hass, entity_id, ATTR_TEMPERATURE) is None
+    assert (state := hass.states.get(entity_id))
+    assert state.attributes[ATTR_TEMPERATURE] is None
 
 
 async def test_blu_trv_climate_hvac_action(
@@ -853,9 +849,11 @@ async def test_blu_trv_climate_hvac_action(
 
     await init_integration(hass, 3, model=MODEL_BLU_GATEWAY_G3)
 
-    assert get_entity_attribute(hass, entity_id, ATTR_HVAC_ACTION) == HVACAction.IDLE
+    assert (state := hass.states.get(entity_id))
+    assert state.attributes[ATTR_HVAC_ACTION] == HVACAction.IDLE
 
     monkeypatch.setitem(mock_blu_trv.status[f"{BLU_TRV_IDENTIFIER}:200"], "pos", 10)
     mock_blu_trv.mock_update()
 
-    assert get_entity_attribute(hass, entity_id, ATTR_HVAC_ACTION) == HVACAction.HEATING
+    assert (state := hass.states.get(entity_id))
+    assert state.attributes[ATTR_HVAC_ACTION] == HVACAction.HEATING

tests/components/shelly/test_config_flow.py

@@ -1080,7 +1080,7 @@ async def test_options_flow_ble(hass: HomeAssistant, mock_rpc_device: Mock) -> N
     await hass.async_block_till_done()
 
     assert result["type"] is FlowResultType.CREATE_ENTRY
-    assert result["data"][CONF_BLE_SCANNER_MODE] == BLEScannerMode.DISABLED
+    assert result["data"][CONF_BLE_SCANNER_MODE] is BLEScannerMode.DISABLED
 
     result = await hass.config_entries.options.async_init(entry.entry_id)
     assert result["type"] is FlowResultType.FORM
@@ -1096,7 +1096,7 @@ async def test_options_flow_ble(hass: HomeAssistant, mock_rpc_device: Mock) -> N
     await hass.async_block_till_done()
 
     assert result["type"] is FlowResultType.CREATE_ENTRY
-    assert result["data"][CONF_BLE_SCANNER_MODE] == BLEScannerMode.ACTIVE
+    assert result["data"][CONF_BLE_SCANNER_MODE] is BLEScannerMode.ACTIVE
 
     result = await hass.config_entries.options.async_init(entry.entry_id)
     assert result["type"] is FlowResultType.FORM
@@ -1112,7 +1112,7 @@ async def test_options_flow_ble(hass: HomeAssistant, mock_rpc_device: Mock) -> N
     await hass.async_block_till_done()
 
     assert result["type"] is FlowResultType.CREATE_ENTRY
-    assert result["data"][CONF_BLE_SCANNER_MODE] == BLEScannerMode.PASSIVE
+    assert result["data"][CONF_BLE_SCANNER_MODE] is BLEScannerMode.PASSIVE
 
     await hass.config_entries.async_unload(entry.entry_id)
 

tests/components/shelly/test_coordinator.py

@@ -32,7 +32,6 @@ from homeassistant.helpers import device_registry as dr, issue_registry as ir
 
 from . import (
     MOCK_MAC,
-    get_entity_state,
     init_integration,
     inject_rpc_device_event,
     mock_polling_rpc_update,
@@ -72,7 +71,7 @@ async def test_block_reload_on_cfg_change(
     async_fire_time_changed(hass)
     await hass.async_block_till_done()
 
-    assert hass.states.get("switch.test_name_channel_1") is not None
+    assert hass.states.get("switch.test_name_channel_1")
 
     # Generate config change from switch to light
     monkeypatch.setitem(
@@ -82,7 +81,7 @@ async def test_block_reload_on_cfg_change(
     mock_block_device.mock_update()
     await hass.async_block_till_done()
 
-    assert hass.states.get("switch.test_name_channel_1") is not None
+    assert hass.states.get("switch.test_name_channel_1")
 
     # Wait for debouncer
     freezer.tick(timedelta(seconds=ENTRY_RELOAD_COOLDOWN))
@@ -114,14 +113,14 @@ async def test_block_no_reload_on_bulb_changes(
     mock_block_device.mock_update()
     await hass.async_block_till_done()
 
-    assert hass.states.get("switch.test_name_channel_1") is not None
+    assert hass.states.get("switch.test_name_channel_1")
 
     # Wait for debouncer
     freezer.tick(timedelta(seconds=ENTRY_RELOAD_COOLDOWN))
     async_fire_time_changed(hass)
     await hass.async_block_till_done()
 
-    assert hass.states.get("switch.test_name_channel_1") is not None
+    assert hass.states.get("switch.test_name_channel_1")
 
     # Test no reload  on effect change
     monkeypatch.setattr(mock_block_device.blocks[LIGHT_BLOCK_ID], "effect", 1)
@@ -129,14 +128,14 @@ async def test_block_no_reload_on_bulb_changes(
     mock_block_device.mock_update()
     await hass.async_block_till_done()
 
-    assert hass.states.get("switch.test_name_channel_1") is not None
+    assert hass.states.get("switch.test_name_channel_1")
 
     # Wait for debouncer
     freezer.tick(timedelta(seconds=ENTRY_RELOAD_COOLDOWN))
     async_fire_time_changed(hass)
     await hass.async_block_till_done()
 
-    assert hass.states.get("switch.test_name_channel_1") is not None
+    assert hass.states.get("switch.test_name_channel_1")
 
 
 async def test_block_polling_auth_error(
@@ -245,14 +244,16 @@ async def test_block_polling_connection_error(
     )
     await init_integration(hass, 1)
 
-    assert get_entity_state(hass, "switch.test_name_channel_1") == STATE_ON
+    assert (state := hass.states.get("switch.test_name_channel_1"))
+    assert state.state == STATE_ON
 
     # Move time to generate polling
     freezer.tick(timedelta(seconds=UPDATE_PERIOD_MULTIPLIER * 15))
     async_fire_time_changed(hass)
     await hass.async_block_till_done()
 
-    assert get_entity_state(hass, "switch.test_name_channel_1") == STATE_UNAVAILABLE
+    assert (state := hass.states.get("switch.test_name_channel_1"))
+    assert state.state == STATE_UNAVAILABLE
 
 
 @pytest.mark.parametrize("exc", [DeviceConnectionError, MacAddressMismatchError])
@@ -270,12 +271,14 @@ async def test_block_rest_update_connection_error(
     await init_integration(hass, 1)
 
     await mock_rest_update(hass, freezer)
-    assert get_entity_state(hass, entity_id) == STATE_ON
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_ON
 
     monkeypatch.setattr(mock_block_device, "update_shelly", AsyncMock(side_effect=exc))
     await mock_rest_update(hass, freezer)
 
-    assert get_entity_state(hass, entity_id) == STATE_UNAVAILABLE
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_UNAVAILABLE
 
 
 async def test_block_sleeping_device_no_periodic_updates(
@@ -297,14 +300,16 @@ async def test_block_sleeping_device_no_periodic_updates(
     mock_block_device.mock_online()
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    assert get_entity_state(hass, entity_id) == "22.1"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "22.1"
 
     # Move time to generate polling
     freezer.tick(timedelta(seconds=UPDATE_PERIOD_MULTIPLIER * 3600))
     async_fire_time_changed(hass)
     await hass.async_block_till_done()
 
-    assert get_entity_state(hass, entity_id) == STATE_UNAVAILABLE
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_UNAVAILABLE
 
 
 async def test_block_device_push_updates_failure(
@@ -416,7 +421,7 @@ async def test_rpc_reload_on_cfg_change(
     )
     await hass.async_block_till_done()
 
-    assert hass.states.get("switch.test_switch_0") is not None
+    assert hass.states.get("switch.test_switch_0")
 
     # Wait for debouncer
     freezer.tick(timedelta(seconds=ENTRY_RELOAD_COOLDOWN))
@@ -596,14 +601,16 @@ async def test_rpc_sleeping_device_no_periodic_updates(
     mock_rpc_device.mock_online()
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    assert get_entity_state(hass, entity_id) == "22.9"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "22.9"
 
     # Move time to generate polling
     freezer.tick(timedelta(seconds=UPDATE_PERIOD_MULTIPLIER * 1000))
     async_fire_time_changed(hass)
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    assert get_entity_state(hass, entity_id) is STATE_UNAVAILABLE
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_UNAVAILABLE
 
 
 async def test_rpc_sleeping_device_firmware_unsupported(
@@ -716,7 +723,8 @@ async def test_rpc_reconnect_error(
     monkeypatch.setitem(mock_rpc_device.status["sys"], "relay_in_thermostat", False)
     await init_integration(hass, 2)
 
-    assert get_entity_state(hass, "switch.test_switch_0") == STATE_ON
+    assert (state := hass.states.get("switch.test_switch_0"))
+    assert state.state == STATE_ON
 
     monkeypatch.setattr(mock_rpc_device, "connected", False)
     monkeypatch.setattr(mock_rpc_device, "initialize", AsyncMock(side_effect=exc))
@@ -726,7 +734,8 @@ async def test_rpc_reconnect_error(
     async_fire_time_changed(hass)
     await hass.async_block_till_done()
 
-    assert get_entity_state(hass, "switch.test_switch_0") == STATE_UNAVAILABLE
+    assert (state := hass.states.get("switch.test_switch_0"))
+    assert state.state == STATE_UNAVAILABLE
 
 
 async def test_rpc_error_running_connected_events(
@@ -748,14 +757,17 @@ async def test_rpc_error_running_connected_events(
         )
 
     assert "Error running connected events for device" in caplog.text
-    assert get_entity_state(hass, "switch.test_switch_0") == STATE_UNAVAILABLE
+
+    assert (state := hass.states.get("switch.test_switch_0"))
+    assert state.state == STATE_UNAVAILABLE
 
     # Move time to generate reconnect without error
     freezer.tick(timedelta(seconds=RPC_RECONNECT_INTERVAL))
     async_fire_time_changed(hass)
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    assert get_entity_state(hass, "switch.test_switch_0") == STATE_ON
+    assert (state := hass.states.get("switch.test_switch_0"))
+    assert state.state == STATE_ON
 
 
 async def test_rpc_polling_connection_error(
@@ -776,11 +788,13 @@ async def test_rpc_polling_connection_error(
         ),
     )
 
-    assert get_entity_state(hass, entity_id) == "-63"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "-63"
 
     await mock_polling_rpc_update(hass, freezer)
 
-    assert get_entity_state(hass, entity_id) == STATE_UNAVAILABLE
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_UNAVAILABLE
 
 
 async def test_rpc_polling_disconnected(
@@ -795,11 +809,13 @@ async def test_rpc_polling_disconnected(
 
     monkeypatch.setattr(mock_rpc_device, "connected", False)
 
-    assert get_entity_state(hass, entity_id) == "-63"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "-63"
 
     await mock_polling_rpc_update(hass, freezer)
 
-    assert get_entity_state(hass, entity_id) == STATE_UNAVAILABLE
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_UNAVAILABLE
 
 
 async def test_rpc_update_entry_fw_ver(
@@ -903,7 +919,8 @@ async def test_block_sleeping_device_connection_error(
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
-    assert get_entity_state(hass, entity_id) == STATE_ON
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_ON
 
     # Make device online event with connection error
     monkeypatch.setattr(
@@ -917,7 +934,8 @@ async def test_block_sleeping_device_connection_error(
     await hass.async_block_till_done(wait_background_tasks=True)
 
     assert "Error connecting to Shelly device" in caplog.text
-    assert get_entity_state(hass, entity_id) == STATE_ON
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_ON
 
     # Move time to generate sleep period update
     freezer.tick(timedelta(seconds=sleep_period * UPDATE_PERIOD_MULTIPLIER))
@@ -925,7 +943,8 @@ async def test_block_sleeping_device_connection_error(
     await hass.async_block_till_done(wait_background_tasks=True)
 
     assert "Sleeping device did not update" in caplog.text
-    assert get_entity_state(hass, entity_id) == STATE_UNAVAILABLE
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_UNAVAILABLE
 
 
 async def test_rpc_sleeping_device_connection_error(
@@ -954,7 +973,8 @@ async def test_rpc_sleeping_device_connection_error(
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
-    assert get_entity_state(hass, entity_id) == STATE_ON
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_ON
 
     # Make device online event with connection error
     monkeypatch.setattr(
@@ -968,7 +988,8 @@ async def test_rpc_sleeping_device_connection_error(
     await hass.async_block_till_done(wait_background_tasks=True)
 
     assert "Error connecting to Shelly device" in caplog.text
-    assert get_entity_state(hass, entity_id) == STATE_ON
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_ON
 
     # Move time to generate sleep period update
     freezer.tick(timedelta(seconds=sleep_period * UPDATE_PERIOD_MULTIPLIER))
@@ -976,7 +997,8 @@ async def test_rpc_sleeping_device_connection_error(
     await hass.async_block_till_done(wait_background_tasks=True)
 
     assert "Sleeping device did not update" in caplog.text
-    assert get_entity_state(hass, entity_id) == STATE_UNAVAILABLE
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_UNAVAILABLE
 
 
 async def test_rpc_sleeping_device_late_setup(
@@ -1001,7 +1023,8 @@ async def test_rpc_sleeping_device_late_setup(
     monkeypatch.setattr(mock_rpc_device, "connected", True)
     mock_rpc_device.mock_initialized()
     await hass.async_block_till_done(wait_background_tasks=True)
-    assert hass.states.get("sensor.test_name_temperature") is not None
+
+    assert hass.states.get("sensor.test_name_temperature")
 
 
 async def test_rpc_already_connected(

tests/components/shelly/test_cover.py

@@ -47,7 +47,7 @@ async def test_block_device_services(
         {ATTR_ENTITY_ID: entity_id, ATTR_POSITION: 50},
         blocking=True,
     )
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.attributes[ATTR_CURRENT_POSITION] == 50
 
     await hass.services.async_call(
@@ -56,7 +56,8 @@ async def test_block_device_services(
         {ATTR_ENTITY_ID: entity_id},
         blocking=True,
     )
-    assert hass.states.get(entity_id).state == CoverState.OPENING
+    assert (state := hass.states.get(entity_id))
+    assert state.state == CoverState.OPENING
 
     await hass.services.async_call(
         COVER_DOMAIN,
@@ -64,7 +65,8 @@ async def test_block_device_services(
         {ATTR_ENTITY_ID: entity_id},
         blocking=True,
     )
-    assert hass.states.get(entity_id).state == CoverState.CLOSING
+    assert (state := hass.states.get(entity_id))
+    assert state.state == CoverState.CLOSING
 
     await hass.services.async_call(
         COVER_DOMAIN,
@@ -72,10 +74,10 @@ async def test_block_device_services(
         {ATTR_ENTITY_ID: entity_id},
         blocking=True,
     )
-    assert hass.states.get(entity_id).state == CoverState.CLOSED
+    assert (state := hass.states.get(entity_id))
+    assert state.state == CoverState.CLOSED
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-roller_0"
 
 
@@ -86,11 +88,15 @@ async def test_block_device_update(
     monkeypatch.setattr(mock_block_device.blocks[ROLLER_BLOCK_ID], "rollerPos", 0)
     await init_integration(hass, 1)
 
-    assert hass.states.get("cover.test_name").state == CoverState.CLOSED
+    state = hass.states.get("cover.test_name")
+    assert state
+    assert state.state == CoverState.CLOSED
 
     monkeypatch.setattr(mock_block_device.blocks[ROLLER_BLOCK_ID], "rollerPos", 100)
     mock_block_device.mock_update()
-    assert hass.states.get("cover.test_name").state == CoverState.OPEN
+    state = hass.states.get("cover.test_name")
+    assert state
+    assert state.state == CoverState.OPEN
 
 
 async def test_block_device_no_roller_blocks(
@@ -99,6 +105,7 @@ async def test_block_device_no_roller_blocks(
     """Test block device without roller blocks."""
     monkeypatch.setattr(mock_block_device.blocks[ROLLER_BLOCK_ID], "type", None)
     await init_integration(hass, 1)
+
     assert hass.states.get("cover.test_name") is None
 
 
@@ -118,7 +125,7 @@ async def test_rpc_device_services(
         {ATTR_ENTITY_ID: entity_id, ATTR_POSITION: 50},
         blocking=True,
     )
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.attributes[ATTR_CURRENT_POSITION] == 50
 
     mutate_rpc_device_status(
@@ -131,7 +138,9 @@ async def test_rpc_device_services(
         blocking=True,
     )
     mock_rpc_device.mock_update()
-    assert hass.states.get(entity_id).state == CoverState.OPENING
+
+    assert (state := hass.states.get(entity_id))
+    assert state.state == CoverState.OPENING
 
     mutate_rpc_device_status(
         monkeypatch, mock_rpc_device, "cover:0", "state", "closing"
@@ -143,7 +152,9 @@ async def test_rpc_device_services(
         blocking=True,
     )
     mock_rpc_device.mock_update()
-    assert hass.states.get(entity_id).state == CoverState.CLOSING
+
+    assert (state := hass.states.get(entity_id))
+    assert state.state == CoverState.CLOSING
 
     mutate_rpc_device_status(monkeypatch, mock_rpc_device, "cover:0", "state", "closed")
     await hass.services.async_call(
@@ -153,10 +164,10 @@ async def test_rpc_device_services(
         blocking=True,
     )
     mock_rpc_device.mock_update()
-    assert hass.states.get(entity_id).state == CoverState.CLOSED
+    assert (state := hass.states.get(entity_id))
+    assert state.state == CoverState.CLOSED
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-cover:0"
 
 
@@ -166,6 +177,7 @@ async def test_rpc_device_no_cover_keys(
     """Test RPC device without cover keys."""
     monkeypatch.delitem(mock_rpc_device.status, "cover:0")
     await init_integration(hass, 2)
+
     assert hass.states.get("cover.test_cover_0") is None
 
 
@@ -175,11 +187,16 @@ async def test_rpc_device_update(
     """Test RPC device update."""
     mutate_rpc_device_status(monkeypatch, mock_rpc_device, "cover:0", "state", "closed")
     await init_integration(hass, 2)
-    assert hass.states.get("cover.test_cover_0").state == CoverState.CLOSED
+
+    state = hass.states.get("cover.test_cover_0")
+    assert state
+    assert state.state == CoverState.CLOSED
 
     mutate_rpc_device_status(monkeypatch, mock_rpc_device, "cover:0", "state", "open")
     mock_rpc_device.mock_update()
-    assert hass.states.get("cover.test_cover_0").state == CoverState.OPEN
+    state = hass.states.get("cover.test_cover_0")
+    assert state
+    assert state.state == CoverState.OPEN
 
 
 async def test_rpc_device_no_position_control(
@@ -190,7 +207,10 @@ async def test_rpc_device_no_position_control(
         monkeypatch, mock_rpc_device, "cover:0", "pos_control", False
     )
     await init_integration(hass, 2)
-    assert hass.states.get("cover.test_cover_0").state == CoverState.OPEN
+
+    state = hass.states.get("cover.test_cover_0")
+    assert state
+    assert state.state == CoverState.OPEN
 
 
 async def test_rpc_cover_tilt(
@@ -212,11 +232,10 @@ async def test_rpc_cover_tilt(
 
     await init_integration(hass, 3)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.attributes[ATTR_CURRENT_TILT_POSITION] == 0
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-cover:0"
 
     await hass.services.async_call(
@@ -228,7 +247,7 @@ async def test_rpc_cover_tilt(
     mutate_rpc_device_status(monkeypatch, mock_rpc_device, "cover:0", "slat_pos", 50)
     mock_rpc_device.mock_update()
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.attributes[ATTR_CURRENT_TILT_POSITION] == 50
 
     await hass.services.async_call(
@@ -240,7 +259,7 @@ async def test_rpc_cover_tilt(
     mutate_rpc_device_status(monkeypatch, mock_rpc_device, "cover:0", "slat_pos", 100)
     mock_rpc_device.mock_update()
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.attributes[ATTR_CURRENT_TILT_POSITION] == 100
 
     await hass.services.async_call(
@@ -258,5 +277,5 @@ async def test_rpc_cover_tilt(
     mutate_rpc_device_status(monkeypatch, mock_rpc_device, "cover:0", "slat_pos", 10)
     mock_rpc_device.mock_update()
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.attributes[ATTR_CURRENT_TILT_POSITION] == 10

tests/components/shelly/test_event.py

@@ -33,8 +33,7 @@ async def test_rpc_button(
     await init_integration(hass, 2)
     entity_id = "event.test_name_input_0"
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert state
     assert state.state == STATE_UNKNOWN
     assert state.attributes.get(ATTR_EVENT_TYPES) == unordered(
         ["btn_down", "btn_up", "double_push", "long_push", "single_push", "triple_push"]
@@ -42,8 +41,7 @@ async def test_rpc_button(
     assert state.attributes.get(ATTR_EVENT_TYPE) is None
     assert state.attributes.get(ATTR_DEVICE_CLASS) == EventDeviceClass.BUTTON
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-input:0"
 
     inject_rpc_device_event(
@@ -62,7 +60,7 @@ async def test_rpc_button(
     )
     await hass.async_block_till_done()
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.attributes.get(ATTR_EVENT_TYPE) == "single_push"
 
 
@@ -78,11 +76,9 @@ async def test_rpc_script_1_event(
     await init_integration(hass, 2)
     entity_id = "event.test_name_test_script_js"
 
-    state = hass.states.get(entity_id)
+    assert hass.states.get(entity_id) == snapshot(name=f"{entity_id}-state")
-    assert state == snapshot(name=f"{entity_id}-state")
 
-    entry = entity_registry.async_get(entity_id)
+    assert entity_registry.async_get(entity_id) == snapshot(name=f"{entity_id}-entry")
-    assert entry == snapshot(name=f"{entity_id}-entry")
 
     inject_rpc_device_event(
         monkeypatch,
@@ -101,7 +97,7 @@ async def test_rpc_script_1_event(
     )
     await hass.async_block_till_done()
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.attributes.get(ATTR_EVENT_TYPE) == "script_start"
 
     inject_rpc_device_event(
@@ -121,7 +117,7 @@ async def test_rpc_script_1_event(
     )
     await hass.async_block_till_done()
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.attributes.get(ATTR_EVENT_TYPE) != "unknown_event"
 
 
@@ -135,11 +131,9 @@ async def test_rpc_script_2_event(
     await init_integration(hass, 2)
     entity_id = "event.test_name_test_script_2_js"
 
-    state = hass.states.get(entity_id)
+    assert hass.states.get(entity_id) == snapshot(name=f"{entity_id}-state")
-    assert state == snapshot(name=f"{entity_id}-state")
 
-    entry = entity_registry.async_get(entity_id)
+    assert entity_registry.async_get(entity_id) == snapshot(name=f"{entity_id}-entry")
-    assert entry == snapshot(name=f"{entity_id}-entry")
 
 
 @pytest.mark.usefixtures("entity_registry_enabled_by_default")
@@ -152,11 +146,9 @@ async def test_rpc_script_ble_event(
     await init_integration(hass, 2)
     entity_id = f"event.test_name_{BLE_SCRIPT_NAME}"
 
-    state = hass.states.get(entity_id)
+    assert hass.states.get(entity_id) == snapshot(name=f"{entity_id}-state")
-    assert state == snapshot(name=f"{entity_id}-state")
 
-    entry = entity_registry.async_get(entity_id)
+    assert entity_registry.async_get(entity_id) == snapshot(name=f"{entity_id}-entry")
-    assert entry == snapshot(name=f"{entity_id}-entry")
 
 
 async def test_rpc_event_removal(
@@ -186,15 +178,13 @@ async def test_block_event(
     await init_integration(hass, 1)
     entity_id = "event.test_name_channel_1"
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert state
     assert state.state == STATE_UNKNOWN
     assert state.attributes.get(ATTR_EVENT_TYPES) == unordered(["single", "long"])
     assert state.attributes.get(ATTR_EVENT_TYPE) is None
     assert state.attributes.get(ATTR_DEVICE_CLASS) == EventDeviceClass.BUTTON
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-relay_0-1"
 
     monkeypatch.setattr(
@@ -206,7 +196,7 @@ async def test_block_event(
     mock_block_device.mock_update()
     await hass.async_block_till_done()
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.attributes.get(ATTR_EVENT_TYPE) == "long"
 
 
@@ -217,8 +207,7 @@ async def test_block_event_shix3_1(
     await init_integration(hass, 1, model=MODEL_I3)
     entity_id = "event.test_name_channel_1"
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert state
     assert state.attributes.get(ATTR_EVENT_TYPES) == unordered(
         ["double", "long", "long_single", "single", "single_long", "triple"]
     )

tests/components/shelly/test_init.py

@@ -307,7 +307,8 @@ async def test_sleeping_rpc_device_online_during_setup(
 
     assert "will resume when device is online" in caplog.text
     assert "is online (source: setup)" in caplog.text
-    assert hass.states.get("sensor.test_name_temperature") is not None
+
+    assert hass.states.get("sensor.test_name_temperature")
 
 
 async def test_sleeping_rpc_device_offline_during_setup(
@@ -336,7 +337,7 @@ async def test_sleeping_rpc_device_offline_during_setup(
     mock_rpc_device.mock_online()
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    assert hass.states.get("sensor.test_name_temperature") is not None
+    assert hass.states.get("sensor.test_name_temperature")
 
 
 @pytest.mark.parametrize(
@@ -360,13 +361,15 @@ async def test_entry_unload(
     entry = await init_integration(hass, gen)
 
     assert entry.state is ConfigEntryState.LOADED
-    assert hass.states.get(entity_id).state is STATE_ON
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_ON
 
     await hass.config_entries.async_unload(entry.entry_id)
     await hass.async_block_till_done()
 
     assert entry.state is ConfigEntryState.NOT_LOADED
-    assert hass.states.get(entity_id).state is STATE_UNAVAILABLE
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_UNAVAILABLE
 
 
 @pytest.mark.parametrize(
@@ -384,9 +387,9 @@ async def test_entry_unload_device_not_ready(
     mock_rpc_device: Mock,
 ) -> None:
     """Test entry unload when device is not ready."""
-    entry = await init_integration(hass, gen, sleep_period=1000)
+    assert (entry := await init_integration(hass, gen, sleep_period=1000))
-
     assert entry.state is ConfigEntryState.LOADED
+
     assert hass.states.get(entity_id) is None
 
     await hass.config_entries.async_unload(entry.entry_id)
@@ -405,13 +408,15 @@ async def test_entry_unload_not_connected(
     with patch(
         "homeassistant.components.shelly.coordinator.async_stop_scanner"
     ) as mock_stop_scanner:
-        entry = await init_integration(
+        assert (
-            hass, 2, options={CONF_BLE_SCANNER_MODE: BLEScannerMode.ACTIVE}
+            entry := await init_integration(
+                hass, 2, options={CONF_BLE_SCANNER_MODE: BLEScannerMode.ACTIVE}
+            )
         )
-        entity_id = "switch.test_switch_0"
-
         assert entry.state is ConfigEntryState.LOADED
-        assert hass.states.get(entity_id).state is STATE_ON
+
+        assert (state := hass.states.get("switch.test_switch_0"))
+        assert state.state == STATE_ON
         assert not mock_stop_scanner.call_count
 
         monkeypatch.setattr(mock_rpc_device, "connected", False)
@@ -434,13 +439,15 @@ async def test_entry_unload_not_connected_but_we_think_we_are(
         "homeassistant.components.shelly.coordinator.async_stop_scanner",
         side_effect=DeviceConnectionError,
     ) as mock_stop_scanner:
-        entry = await init_integration(
+        assert (
-            hass, 2, options={CONF_BLE_SCANNER_MODE: BLEScannerMode.ACTIVE}
+            entry := await init_integration(
+                hass, 2, options={CONF_BLE_SCANNER_MODE: BLEScannerMode.ACTIVE}
+            )
         )
-        entity_id = "switch.test_switch_0"
-
         assert entry.state is ConfigEntryState.LOADED
-        assert hass.states.get(entity_id).state is STATE_ON
+
+        assert (state := hass.states.get("switch.test_switch_0"))
+        assert state.state == STATE_ON
         assert not mock_stop_scanner.call_count
 
         monkeypatch.setattr(mock_rpc_device, "connected", False)
@@ -473,7 +480,9 @@ async def test_entry_missing_gen(hass: HomeAssistant, mock_block_device: Mock) -
     entry = await init_integration(hass, None)
 
     assert entry.state is ConfigEntryState.LOADED
-    assert hass.states.get("switch.test_name_channel_1").state is STATE_ON
+
+    assert (state := hass.states.get("switch.test_name_channel_1"))
+    assert state.state == STATE_ON
 
 
 async def test_entry_missing_port(hass: HomeAssistant) -> None:

tests/components/shelly/test_light.py

@@ -65,18 +65,17 @@ async def test_block_device_rgbw_bulb(
     await init_integration(hass, 1, model=MODEL_BULB)
 
     # Test initial
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    attributes = state.attributes
     assert state.state == STATE_ON
-    assert attributes[ATTR_RGBW_COLOR] == (45, 55, 65, 70)
+    assert state.attributes[ATTR_RGBW_COLOR] == (45, 55, 65, 70)
-    assert attributes[ATTR_BRIGHTNESS] == 48
+    assert state.attributes[ATTR_BRIGHTNESS] == 48
-    assert attributes[ATTR_SUPPORTED_COLOR_MODES] == [
+    assert state.attributes[ATTR_SUPPORTED_COLOR_MODES] == [
         ColorMode.COLOR_TEMP,
         ColorMode.RGBW,
     ]
-    assert attributes[ATTR_SUPPORTED_FEATURES] == LightEntityFeature.EFFECT
+    assert state.attributes[ATTR_SUPPORTED_FEATURES] == LightEntityFeature.EFFECT
-    assert len(attributes[ATTR_EFFECT_LIST]) == 7
+    assert len(state.attributes[ATTR_EFFECT_LIST]) == 7
-    assert attributes[ATTR_EFFECT] == "Off"
+    assert state.attributes[ATTR_EFFECT] == "Off"
 
     # Turn off
     mock_block_device.blocks[LIGHT_BLOCK_ID].set_state.reset_mock()
@@ -89,7 +88,7 @@ async def test_block_device_rgbw_bulb(
     mock_block_device.blocks[LIGHT_BLOCK_ID].set_state.assert_called_once_with(
         turn="off"
     )
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_OFF
 
     # Turn on, RGBW = [70, 80, 90, 20], brightness = 33, effect = Flash
@@ -108,13 +107,12 @@ async def test_block_device_rgbw_bulb(
     mock_block_device.blocks[LIGHT_BLOCK_ID].set_state.assert_called_once_with(
         turn="on", gain=13, brightness=13, red=70, green=80, blue=90, white=30, effect=3
     )
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    attributes = state.attributes
     assert state.state == STATE_ON
-    assert attributes[ATTR_COLOR_MODE] == ColorMode.RGBW
+    assert state.attributes[ATTR_COLOR_MODE] == ColorMode.RGBW
-    assert attributes[ATTR_RGBW_COLOR] == (70, 80, 90, 30)
+    assert state.attributes[ATTR_RGBW_COLOR] == (70, 80, 90, 30)
-    assert attributes[ATTR_BRIGHTNESS] == 33
+    assert state.attributes[ATTR_BRIGHTNESS] == 33
-    assert attributes[ATTR_EFFECT] == "Flash"
+    assert state.attributes[ATTR_EFFECT] == "Flash"
 
     # Turn on, COLOR_TEMP_KELVIN = 3500
     mock_block_device.blocks[LIGHT_BLOCK_ID].set_state.reset_mock()
@@ -127,14 +125,12 @@ async def test_block_device_rgbw_bulb(
     mock_block_device.blocks[LIGHT_BLOCK_ID].set_state.assert_called_once_with(
         turn="on", temp=3500, mode="white"
     )
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    attributes = state.attributes
     assert state.state == STATE_ON
-    assert attributes[ATTR_COLOR_MODE] == ColorMode.COLOR_TEMP
+    assert state.attributes[ATTR_COLOR_MODE] == ColorMode.COLOR_TEMP
-    assert attributes[ATTR_COLOR_TEMP_KELVIN] == 3500
+    assert state.attributes[ATTR_COLOR_TEMP_KELVIN] == 3500
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-light_0"
 
 
@@ -154,21 +150,20 @@ async def test_block_device_rgb_bulb(
     await init_integration(hass, 1, model=MODEL_BULB_RGBW)
 
     # Test initial
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    attributes = state.attributes
     assert state.state == STATE_ON
-    assert attributes[ATTR_RGB_COLOR] == (45, 55, 65)
+    assert state.attributes[ATTR_RGB_COLOR] == (45, 55, 65)
-    assert attributes[ATTR_BRIGHTNESS] == 48
+    assert state.attributes[ATTR_BRIGHTNESS] == 48
-    assert attributes[ATTR_SUPPORTED_COLOR_MODES] == [
+    assert state.attributes[ATTR_SUPPORTED_COLOR_MODES] == [
         ColorMode.COLOR_TEMP,
         ColorMode.RGB,
     ]
     assert (
-        attributes[ATTR_SUPPORTED_FEATURES]
+        state.attributes[ATTR_SUPPORTED_FEATURES]
         == LightEntityFeature.EFFECT | LightEntityFeature.TRANSITION
     )
-    assert len(attributes[ATTR_EFFECT_LIST]) == 4
+    assert len(state.attributes[ATTR_EFFECT_LIST]) == 4
-    assert attributes[ATTR_EFFECT] == "Off"
+    assert state.attributes[ATTR_EFFECT] == "Off"
 
     # Turn off
     mock_block_device.blocks[LIGHT_BLOCK_ID].set_state.reset_mock()
@@ -181,7 +176,7 @@ async def test_block_device_rgb_bulb(
     mock_block_device.blocks[LIGHT_BLOCK_ID].set_state.assert_called_once_with(
         turn="off"
     )
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_OFF
 
     # Turn on, RGB = [70, 80, 90], brightness = 33, effect = Flash
@@ -200,13 +195,12 @@ async def test_block_device_rgb_bulb(
     mock_block_device.blocks[LIGHT_BLOCK_ID].set_state.assert_called_once_with(
         turn="on", gain=13, brightness=13, red=70, green=80, blue=90, effect=3
     )
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    attributes = state.attributes
     assert state.state == STATE_ON
-    assert attributes[ATTR_COLOR_MODE] == ColorMode.RGB
+    assert state.attributes[ATTR_COLOR_MODE] == ColorMode.RGB
-    assert attributes[ATTR_RGB_COLOR] == (70, 80, 90)
+    assert state.attributes[ATTR_RGB_COLOR] == (70, 80, 90)
-    assert attributes[ATTR_BRIGHTNESS] == 33
+    assert state.attributes[ATTR_BRIGHTNESS] == 33
-    assert attributes[ATTR_EFFECT] == "Flash"
+    assert state.attributes[ATTR_EFFECT] == "Flash"
 
     # Turn on, COLOR_TEMP_KELVIN = 3500
     mock_block_device.blocks[LIGHT_BLOCK_ID].set_state.reset_mock()
@@ -219,11 +213,10 @@ async def test_block_device_rgb_bulb(
     mock_block_device.blocks[LIGHT_BLOCK_ID].set_state.assert_called_once_with(
         turn="on", temp=3500, mode="white"
     )
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    attributes = state.attributes
     assert state.state == STATE_ON
-    assert attributes[ATTR_COLOR_MODE] == ColorMode.COLOR_TEMP
+    assert state.attributes[ATTR_COLOR_MODE] == ColorMode.COLOR_TEMP
-    assert attributes[ATTR_COLOR_TEMP_KELVIN] == 3500
+    assert state.attributes[ATTR_COLOR_TEMP_KELVIN] == 3500
 
     # Turn on with unsupported effect
     mock_block_device.blocks[LIGHT_BLOCK_ID].set_state.reset_mock()
@@ -236,14 +229,13 @@ async def test_block_device_rgb_bulb(
     mock_block_device.blocks[LIGHT_BLOCK_ID].set_state.assert_called_once_with(
         turn="on", mode="color"
     )
-    state = hass.states.get(entity_id)
+
-    attributes = state.attributes
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_ON
-    assert attributes[ATTR_EFFECT] == "Off"
+    assert state.attributes[ATTR_EFFECT] == "Off"
     assert "Effect 'Breath' not supported" in caplog.text
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-light_1"
 
 
@@ -272,12 +264,11 @@ async def test_block_device_white_bulb(
     await init_integration(hass, 1, model=MODEL_VINTAGE_V2)
 
     # Test initial
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    attributes = state.attributes
     assert state.state == STATE_ON
-    assert attributes[ATTR_BRIGHTNESS] == 128
+    assert state.attributes[ATTR_BRIGHTNESS] == 128
-    assert attributes[ATTR_SUPPORTED_COLOR_MODES] == [ColorMode.BRIGHTNESS]
+    assert state.attributes[ATTR_SUPPORTED_COLOR_MODES] == [ColorMode.BRIGHTNESS]
-    assert attributes[ATTR_SUPPORTED_FEATURES] == LightEntityFeature.TRANSITION
+    assert state.attributes[ATTR_SUPPORTED_FEATURES] == LightEntityFeature.TRANSITION
 
     # Turn off
     mock_block_device.blocks[LIGHT_BLOCK_ID].set_state.reset_mock()
@@ -290,7 +281,7 @@ async def test_block_device_white_bulb(
     mock_block_device.blocks[LIGHT_BLOCK_ID].set_state.assert_called_once_with(
         turn="off"
     )
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_OFF
 
     # Turn on, brightness = 33
@@ -304,13 +295,11 @@ async def test_block_device_white_bulb(
     mock_block_device.blocks[LIGHT_BLOCK_ID].set_state.assert_called_once_with(
         turn="on", gain=13, brightness=13
     )
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    attributes = state.attributes
     assert state.state == STATE_ON
-    assert attributes[ATTR_BRIGHTNESS] == 33
+    assert state.attributes[ATTR_BRIGHTNESS] == 33
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-light_1"
 
 
@@ -343,9 +332,8 @@ async def test_block_device_support_transition(
     await init_integration(hass, 1, model=model)
 
     # Test initial
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    attributes = state.attributes
+    assert state.attributes[ATTR_SUPPORTED_FEATURES] & LightEntityFeature.TRANSITION
-    assert attributes[ATTR_SUPPORTED_FEATURES] & LightEntityFeature.TRANSITION
 
     # Turn on, TRANSITION = 4
     mock_block_device.blocks[LIGHT_BLOCK_ID].set_state.reset_mock()
@@ -358,7 +346,7 @@ async def test_block_device_support_transition(
     mock_block_device.blocks[LIGHT_BLOCK_ID].set_state.assert_called_once_with(
         turn="on", transition=4000
     )
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_ON
 
     # Turn off, TRANSITION = 6, limit to 5000ms
@@ -372,11 +360,10 @@ async def test_block_device_support_transition(
     mock_block_device.blocks[LIGHT_BLOCK_ID].set_state.assert_called_once_with(
         turn="off", transition=5000
     )
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_OFF
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-light_1"
 
 
@@ -403,14 +390,14 @@ async def test_block_device_relay_app_type_light(
         mock_block_device.blocks[RELAY_BLOCK_ID], "description", "relay_1"
     )
     await init_integration(hass, 1)
+
     assert hass.states.get("switch.test_name_channel_1") is None
 
     # Test initial
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    attributes = state.attributes
     assert state.state == STATE_ON
-    assert attributes[ATTR_SUPPORTED_COLOR_MODES] == [ColorMode.ONOFF]
+    assert state.attributes[ATTR_SUPPORTED_COLOR_MODES] == [ColorMode.ONOFF]
-    assert attributes[ATTR_SUPPORTED_FEATURES] == 0
+    assert state.attributes[ATTR_SUPPORTED_FEATURES] == 0
 
     # Turn off
     mock_block_device.blocks[RELAY_BLOCK_ID].set_state.reset_mock()
@@ -423,7 +410,7 @@ async def test_block_device_relay_app_type_light(
     mock_block_device.blocks[RELAY_BLOCK_ID].set_state.assert_called_once_with(
         turn="off"
     )
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_OFF
 
     # Turn on
@@ -437,11 +424,10 @@ async def test_block_device_relay_app_type_light(
     mock_block_device.blocks[RELAY_BLOCK_ID].set_state.assert_called_once_with(
         turn="on"
     )
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_ON
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-relay_1"
 
 
@@ -451,6 +437,7 @@ async def test_block_device_no_light_blocks(
     """Test block device without light blocks."""
     monkeypatch.setattr(mock_block_device.blocks[LIGHT_BLOCK_ID], "type", "roller")
     await init_integration(hass, 1)
+
     assert hass.states.get("light.test_name_channel_1") is None
 
 
@@ -473,7 +460,9 @@ async def test_rpc_device_switch_type_lights_mode(
         {ATTR_ENTITY_ID: entity_id},
         blocking=True,
     )
-    assert hass.states.get(entity_id).state == STATE_ON
+
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_ON
 
     mutate_rpc_device_status(monkeypatch, mock_rpc_device, "switch:0", "output", False)
     await hass.services.async_call(
@@ -483,10 +472,11 @@ async def test_rpc_device_switch_type_lights_mode(
         blocking=True,
     )
     mock_rpc_device.mock_update()
-    assert hass.states.get(entity_id).state == STATE_OFF
 
-    entry = entity_registry.async_get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert entry
+    assert state.state == STATE_OFF
+
+    assert (entry := entity_registry.async_get(entity_id))
     assert entry.unique_id == "123456789ABC-switch:0"
 
 
@@ -510,7 +500,8 @@ async def test_rpc_light(
     )
 
     mock_rpc_device.call_rpc.assert_called_once_with("Light.Set", {"id": 0, "on": True})
-    state = hass.states.get(entity_id)
+
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_ON
     assert state.attributes[ATTR_BRIGHTNESS] == 135
 
@@ -528,7 +519,8 @@ async def test_rpc_light(
     mock_rpc_device.call_rpc.assert_called_once_with(
         "Light.Set", {"id": 0, "on": False}
     )
-    state = hass.states.get(entity_id)
+
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_OFF
 
     # Turn on, brightness = 33
@@ -547,7 +539,8 @@ async def test_rpc_light(
     mock_rpc_device.call_rpc.assert_called_once_with(
         "Light.Set", {"id": 0, "on": True, "brightness": 13}
     )
-    state = hass.states.get(entity_id)
+
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_ON
     assert state.attributes[ATTR_BRIGHTNESS] == 33
 
@@ -565,7 +558,8 @@ async def test_rpc_light(
     mock_rpc_device.call_rpc.assert_called_once_with(
         "Light.Set", {"id": 0, "on": True, "transition_duration": 10.1}
     )
-    state = hass.states.get(entity_id)
+
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_ON
 
     # Turn off, transition = 0.4, should be limited to 0.5
@@ -584,11 +578,10 @@ async def test_rpc_light(
         "Light.Set", {"id": 0, "on": False, "transition_duration": 0.5}
     )
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_OFF
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-light:0"
 
 
@@ -606,12 +599,11 @@ async def test_rpc_device_rgb_profile(
     await init_integration(hass, 2)
 
     # Test initial
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    attributes = state.attributes
     assert state.state == STATE_ON
-    assert attributes[ATTR_RGB_COLOR] == (45, 55, 65)
+    assert state.attributes[ATTR_RGB_COLOR] == (45, 55, 65)
-    assert attributes[ATTR_SUPPORTED_COLOR_MODES] == [ColorMode.RGB]
+    assert state.attributes[ATTR_SUPPORTED_COLOR_MODES] == [ColorMode.RGB]
-    assert attributes[ATTR_SUPPORTED_FEATURES] == LightEntityFeature.TRANSITION
+    assert state.attributes[ATTR_SUPPORTED_FEATURES] == LightEntityFeature.TRANSITION
 
     # Turn on, RGB = [70, 80, 90]
     await hass.services.async_call(
@@ -628,14 +620,12 @@ async def test_rpc_device_rgb_profile(
         "RGB.Set", {"id": 0, "on": True, "rgb": [70, 80, 90]}
     )
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    attributes = state.attributes
     assert state.state == STATE_ON
-    assert attributes[ATTR_COLOR_MODE] == ColorMode.RGB
+    assert state.attributes[ATTR_COLOR_MODE] == ColorMode.RGB
-    assert attributes[ATTR_RGB_COLOR] == (70, 80, 90)
+    assert state.attributes[ATTR_RGB_COLOR] == (70, 80, 90)
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-rgb:0"
 
 
@@ -653,12 +643,11 @@ async def test_rpc_device_rgbw_profile(
     await init_integration(hass, 2)
 
     # Test initial
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    attributes = state.attributes
     assert state.state == STATE_ON
-    assert attributes[ATTR_RGBW_COLOR] == (21, 22, 23, 120)
+    assert state.attributes[ATTR_RGBW_COLOR] == (21, 22, 23, 120)
-    assert attributes[ATTR_SUPPORTED_COLOR_MODES] == [ColorMode.RGBW]
+    assert state.attributes[ATTR_SUPPORTED_COLOR_MODES] == [ColorMode.RGBW]
-    assert attributes[ATTR_SUPPORTED_FEATURES] == LightEntityFeature.TRANSITION
+    assert state.attributes[ATTR_SUPPORTED_FEATURES] == LightEntityFeature.TRANSITION
 
     # Turn on, RGBW = [72, 82, 92, 128]
     await hass.services.async_call(
@@ -678,14 +667,12 @@ async def test_rpc_device_rgbw_profile(
         "RGBW.Set", {"id": 0, "on": True, "rgb": [72, 82, 92], "white": 128}
     )
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    attributes = state.attributes
     assert state.state == STATE_ON
-    assert attributes[ATTR_COLOR_MODE] == ColorMode.RGBW
+    assert state.attributes[ATTR_COLOR_MODE] == ColorMode.RGBW
-    assert attributes[ATTR_RGBW_COLOR] == (72, 82, 92, 128)
+    assert state.attributes[ATTR_RGBW_COLOR] == (72, 82, 92, 128)
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-rgbw:0"
 
 
@@ -730,9 +717,11 @@ async def test_rpc_rgbw_device_light_mode_remove_others(
     # verify we have 4 lights
     for i in range(SHELLY_PLUS_RGBW_CHANNELS):
         entity_id = f"light.test_light_{i}"
-        assert hass.states.get(entity_id).state == STATE_ON
+
-        entry = entity_registry.async_get(entity_id)
+        assert (state := hass.states.get(entity_id))
-        assert entry
+        assert state.state == STATE_ON
+
+        assert (entry := entity_registry.async_get(entity_id))
         assert entry.unique_id == f"123456789ABC-light:{i}"
 
     # verify RGB & RGBW entities removed
@@ -793,9 +782,11 @@ async def test_rpc_rgbw_device_rgb_w_modes_remove_others(
 
     # verify we have RGB/w light
     entity_id = f"light.test_{active_mode}_0"
-    assert hass.states.get(entity_id).state == STATE_ON
+
-    entry = entity_registry.async_get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert entry
+    assert state.state == STATE_ON
+
+    assert (entry := entity_registry.async_get(entity_id))
     assert entry.unique_id == f"123456789ABC-{active_mode}:0"
 
     # verify light & RGB/W entities removed
@@ -823,8 +814,7 @@ async def test_rpc_cct_light(
 
     await init_integration(hass, 2)
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-cct:0"
 
     # Turn off
@@ -836,7 +826,8 @@ async def test_rpc_cct_light(
     )
 
     mock_rpc_device.call_rpc.assert_called_once_with("CCT.Set", {"id": 0, "on": False})
-    state = hass.states.get(entity_id)
+
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_OFF
 
     # Turn on
@@ -851,7 +842,8 @@ async def test_rpc_cct_light(
 
     mock_rpc_device.mock_update()
     mock_rpc_device.call_rpc.assert_called_once_with("CCT.Set", {"id": 0, "on": True})
-    state = hass.states.get(entity_id)
+
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_ON
     assert state.attributes[ATTR_COLOR_MODE] == ColorMode.COLOR_TEMP
     assert state.attributes[ATTR_BRIGHTNESS] == 196  # 77% of 255
@@ -874,7 +866,8 @@ async def test_rpc_cct_light(
     mock_rpc_device.call_rpc.assert_called_once_with(
         "CCT.Set", {"id": 0, "on": True, "brightness": 88}
     )
-    state = hass.states.get(entity_id)
+
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_ON
     assert state.attributes[ATTR_BRIGHTNESS] == 224  # 88% of 255
 
@@ -894,7 +887,8 @@ async def test_rpc_cct_light(
     mock_rpc_device.call_rpc.assert_called_once_with(
         "CCT.Set", {"id": 0, "on": True, "ct": 4444}
     )
-    state = hass.states.get(entity_id)
+
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_ON
     assert state.attributes[ATTR_COLOR_TEMP_KELVIN] == 4444
 

tests/components/shelly/test_number.py

@@ -54,15 +54,16 @@ async def test_block_number_update(
     mock_block_device.mock_online()
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    assert hass.states.get(entity_id).state == "50"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "50"
 
     monkeypatch.setattr(mock_block_device.blocks[DEVICE_BLOCK_ID], "valvePos", 30)
     mock_block_device.mock_update()
 
-    assert hass.states.get(entity_id).state == "30"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "30"
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-device_0-valvePos"
 
 
@@ -103,14 +104,16 @@ async def test_block_restored_number(
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
-    assert hass.states.get(entity_id).state == "40"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "40"
 
     # Make device online
     monkeypatch.setattr(mock_block_device, "initialized", True)
     mock_block_device.mock_online()
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    assert hass.states.get(entity_id).state == "50"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "50"
 
 
 async def test_block_restored_number_no_last_state(
@@ -141,14 +144,16 @@ async def test_block_restored_number_no_last_state(
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
-    assert hass.states.get(entity_id).state == STATE_UNKNOWN
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_UNKNOWN
 
     # Make device online
     monkeypatch.setattr(mock_block_device, "initialized", True)
     mock_block_device.mock_online()
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    assert hass.states.get(entity_id).state == "50"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "50"
 
 
 async def test_block_number_set_value(
@@ -305,8 +310,7 @@ async def test_rpc_device_virtual_number(
 
     await init_integration(hass, 3)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert state
     assert state.state == "12.3"
     assert state.attributes.get(ATTR_MIN) == 0
     assert state.attributes.get(ATTR_MAX) == 100
@@ -314,13 +318,13 @@ async def test_rpc_device_virtual_number(
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == expected_unit
     assert state.attributes.get(ATTR_MODE) is mode
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-number:203-number"
 
     monkeypatch.setitem(mock_rpc_device.status["number:203"], "value", 78.9)
     mock_rpc_device.mock_update()
-    assert hass.states.get(entity_id).state == "78.9"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "78.9"
 
     monkeypatch.setitem(mock_rpc_device.status["number:203"], "value", 56.7)
     await hass.services.async_call(
@@ -330,7 +334,8 @@ async def test_rpc_device_virtual_number(
         blocking=True,
     )
     mock_rpc_device.mock_update()
-    assert hass.states.get(entity_id).state == "56.7"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "56.7"
 
 
 async def test_rpc_remove_virtual_number_when_mode_label(
@@ -368,8 +373,7 @@ async def test_rpc_remove_virtual_number_when_mode_label(
     await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
 
-    entry = entity_registry.async_get(entity_id)
+    assert entity_registry.async_get(entity_id) is None
-    assert not entry
 
 
 async def test_rpc_remove_virtual_number_when_orphaned(
@@ -393,8 +397,7 @@ async def test_rpc_remove_virtual_number_when_orphaned(
     await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
 
-    entry = entity_registry.async_get(entity_id)
+    assert entity_registry.async_get(entity_id) is None
-    assert not entry
 
 
 async def test_blu_trv_number_entity(
@@ -430,7 +433,8 @@ async def test_blu_trv_ext_temp_set_value(
 
     # After HA start the state should be unknown because there was no previous external
     # temperature report
-    assert hass.states.get(entity_id).state is STATE_UNKNOWN
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_UNKNOWN
 
     await hass.services.async_call(
         NUMBER_DOMAIN,
@@ -452,7 +456,8 @@ async def test_blu_trv_ext_temp_set_value(
         BLU_TRV_TIMEOUT,
     )
 
-    assert hass.states.get(entity_id).state == "22.2"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "22.2"
 
 
 async def test_blu_trv_valve_pos_set_value(
@@ -468,7 +473,8 @@ async def test_blu_trv_valve_pos_set_value(
 
     entity_id = f"{NUMBER_DOMAIN}.trv_name_valve_position"
 
-    assert hass.states.get(entity_id).state == "0"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "0"
 
     monkeypatch.setitem(mock_blu_trv.status["blutrv:200"], "pos", 20)
     await hass.services.async_call(
@@ -493,4 +499,5 @@ async def test_blu_trv_valve_pos_set_value(
     # device only accepts int for 'pos' value
     assert isinstance(mock_blu_trv.call_rpc.call_args[0][1]["params"]["pos"], int)
 
-    assert hass.states.get(entity_id).state == "20"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "20"

tests/components/shelly/test_select.py

@@ -56,8 +56,7 @@ async def test_rpc_device_virtual_enum(
 
     await init_integration(hass, 3)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert state
     assert state.state == expected_state
     assert state.attributes.get(ATTR_OPTIONS) == [
         "Title 1",
@@ -65,13 +64,14 @@ async def test_rpc_device_virtual_enum(
         "option 3",
     ]
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-enum:203-enum"
 
     monkeypatch.setitem(mock_rpc_device.status["enum:203"], "value", "option 2")
     mock_rpc_device.mock_update()
-    assert hass.states.get(entity_id).state == "option 2"
+
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "option 2"
 
     monkeypatch.setitem(mock_rpc_device.status["enum:203"], "value", "option 1")
     await hass.services.async_call(
@@ -83,7 +83,9 @@ async def test_rpc_device_virtual_enum(
     # 'Title 1' corresponds to 'option 1'
     assert mock_rpc_device.call_rpc.call_args[0][1] == {"id": 203, "value": "option 1"}
     mock_rpc_device.mock_update()
-    assert hass.states.get(entity_id).state == "Title 1"
+
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "Title 1"
 
 
 async def test_rpc_remove_virtual_enum_when_mode_label(
@@ -122,8 +124,7 @@ async def test_rpc_remove_virtual_enum_when_mode_label(
     await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
 
-    entry = entity_registry.async_get(entity_id)
+    assert entity_registry.async_get(entity_id) is None
-    assert not entry
 
 
 async def test_rpc_remove_virtual_enum_when_orphaned(
@@ -147,5 +148,4 @@ async def test_rpc_remove_virtual_enum_when_orphaned(
     await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
 
-    entry = entity_registry.async_get(entity_id)
+    assert entity_registry.async_get(entity_id) is None
-    assert not entry

tests/components/shelly/test_sensor.py

@@ -40,7 +40,6 @@ from homeassistant.helpers.entity_registry import EntityRegistry
 from homeassistant.setup import async_setup_component
 
 from . import (
-    get_entity_state,
     init_integration,
     mock_polling_rpc_update,
     mock_rest_update,
@@ -66,15 +65,16 @@ async def test_block_sensor(
     entity_id = f"{SENSOR_DOMAIN}.test_name_channel_1_power"
     await init_integration(hass, 1)
 
-    assert hass.states.get(entity_id).state == "53.4"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "53.4"
 
     monkeypatch.setattr(mock_block_device.blocks[RELAY_BLOCK_ID], "power", 60.1)
     mock_block_device.mock_update()
 
-    assert hass.states.get(entity_id).state == "60.1"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "60.1"
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-relay_0-power"
 
 
@@ -85,14 +85,13 @@ async def test_energy_sensor(
     entity_id = f"{SENSOR_DOMAIN}.test_name_channel_1_energy"
     await init_integration(hass, 1)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     # 1234567.89 Wmin / 60 / 1000 = 20.5761315 kWh
     assert state.state == "20.5761315"
     # suggested unit is KWh
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == UnitOfEnergy.KILO_WATT_HOUR
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-relay_0-energy"
 
 
@@ -111,13 +110,12 @@ async def test_power_factory_unit_migration(
     entity_id = f"{SENSOR_DOMAIN}.test_name_power_factor"
     await init_integration(hass, 1)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     # Value of 0.98 is converted to 98.0%
     assert state.state == "98.0"
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == PERCENTAGE
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-emeter_0-powerFactor"
 
 
@@ -128,12 +126,11 @@ async def test_power_factory_without_unit_migration(
     entity_id = f"{SENSOR_DOMAIN}.test_name_power_factor"
     await init_integration(hass, 1)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == "0.98"
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) is None
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-emeter_0-powerFactor"
 
 
@@ -147,12 +144,14 @@ async def test_block_rest_sensor(
     entity_id = register_entity(hass, SENSOR_DOMAIN, "test_name_rssi", "rssi")
     await init_integration(hass, 1)
 
-    assert hass.states.get(entity_id).state == "-64"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "-64"
 
     monkeypatch.setitem(mock_block_device.status["wifi_sta"], "rssi", -71)
     await mock_rest_update(hass, freezer)
 
-    assert hass.states.get(entity_id).state == "-71"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "-71"
 
 
 async def test_block_sleeping_sensor(
@@ -175,15 +174,16 @@ async def test_block_sleeping_sensor(
     mock_block_device.mock_online()
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    assert hass.states.get(entity_id).state == "22.1"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "22.1"
 
     monkeypatch.setattr(mock_block_device.blocks[SENSOR_BLOCK_ID], "temp", 23.4)
     mock_block_device.mock_update()
 
-    assert hass.states.get(entity_id).state == "23.4"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "23.4"
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-sensor_0-temp"
 
 
@@ -211,8 +211,7 @@ async def test_block_restored_sleeping_sensor(
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert state
     assert state.state == "20.4"
     assert state.attributes[ATTR_STATE_CLASS] == SensorStateClass.MEASUREMENT
     assert state.attributes[ATTR_DEVICE_CLASS] == SensorDeviceClass.TEMPERATURE
@@ -222,7 +221,8 @@ async def test_block_restored_sleeping_sensor(
     mock_block_device.mock_online()
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    assert hass.states.get(entity_id).state == "22.1"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "22.1"
 
 
 async def test_block_restored_sleeping_sensor_no_last_state(
@@ -246,14 +246,16 @@ async def test_block_restored_sleeping_sensor_no_last_state(
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
-    assert hass.states.get(entity_id).state == STATE_UNKNOWN
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_UNKNOWN
 
     # Make device online
     monkeypatch.setattr(mock_block_device, "initialized", True)
     mock_block_device.mock_online()
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    assert hass.states.get(entity_id).state == "22.1"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "22.1"
 
 
 async def test_block_sensor_error(
@@ -266,15 +268,16 @@ async def test_block_sensor_error(
     entity_id = f"{SENSOR_DOMAIN}.test_name_battery"
     await init_integration(hass, 1)
 
-    assert hass.states.get(entity_id).state == "98"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "98"
 
     monkeypatch.setattr(mock_block_device.blocks[DEVICE_BLOCK_ID], "battery", -1)
     mock_block_device.mock_update()
 
-    assert hass.states.get(entity_id).state == STATE_UNAVAILABLE
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_UNAVAILABLE
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-device_0-battery"
 
 
@@ -321,7 +324,8 @@ async def test_block_not_matched_restored_sleeping_sensor(
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
-    assert hass.states.get(entity_id).state == "20.4"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "20.4"
 
     # Make device online
     monkeypatch.setattr(
@@ -331,7 +335,8 @@ async def test_block_not_matched_restored_sleeping_sensor(
     mock_block_device.mock_online()
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    assert hass.states.get(entity_id).state == "20.4"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "20.4"
 
 
 async def test_block_sensor_without_value(
@@ -403,7 +408,8 @@ async def test_block_sensor_values(
     monkeypatch.setattr(mock_block_device.blocks[block_id], attribute, value)
     mock_block_device.mock_update()
 
-    assert hass.states.get(entity_id).state == final_value
+    assert (state := hass.states.get(entity_id))
+    assert state.state == final_value
 
 
 @pytest.mark.parametrize(
@@ -430,7 +436,8 @@ async def test_block_shelly_air_lamp_life(
     )
     await init_integration(hass, 1)
 
-    assert hass.states.get(entity_id).state == percentage
+    assert (state := hass.states.get(entity_id))
+    assert state.state == percentage
 
 
 async def test_rpc_sensor(
@@ -440,17 +447,20 @@ async def test_rpc_sensor(
     entity_id = f"{SENSOR_DOMAIN}.test_cover_0_power"
     await init_integration(hass, 2)
 
-    assert hass.states.get(entity_id).state == "85.3"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "85.3"
 
     mutate_rpc_device_status(monkeypatch, mock_rpc_device, "cover:0", "apower", "88.2")
     mock_rpc_device.mock_update()
 
-    assert hass.states.get(entity_id).state == "88.2"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "88.2"
 
     mutate_rpc_device_status(monkeypatch, mock_rpc_device, "cover:0", "apower", None)
     mock_rpc_device.mock_update()
 
-    assert hass.states.get(entity_id).state == STATE_UNKNOWN
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_UNKNOWN
 
 
 @pytest.mark.usefixtures("entity_registry_enabled_by_default")
@@ -464,7 +474,8 @@ async def test_rpc_rssi_sensor_removal(
     entry = await init_integration(hass, 2)
 
     # WiFi1 enabled, do not remove sensor
-    assert get_entity_state(hass, entity_id) == "-63"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "-63"
 
     # WiFi1 & WiFi2 disabled - remove sensor
     monkeypatch.setitem(mock_rpc_device.config["wifi"]["sta"], "enable", False)
@@ -476,7 +487,9 @@ async def test_rpc_rssi_sensor_removal(
     monkeypatch.setitem(mock_rpc_device.config["wifi"]["sta1"], "enable", True)
     await hass.config_entries.async_reload(entry.entry_id)
     await hass.async_block_till_done()
-    assert get_entity_state(hass, entity_id) == "-63"
+
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "-63"
 
 
 async def test_rpc_illuminance_sensor(
@@ -486,10 +499,10 @@ async def test_rpc_illuminance_sensor(
     entity_id = f"{SENSOR_DOMAIN}.test_name_illuminance"
     await init_integration(hass, 2)
 
-    assert hass.states.get(entity_id).state == "345"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "345"
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-illuminance:0-illuminance"
 
 
@@ -503,17 +516,18 @@ async def test_rpc_sensor_error(
     entity_id = f"{SENSOR_DOMAIN}.test_name_voltmeter"
     await init_integration(hass, 2)
 
-    assert hass.states.get(entity_id).state == "4.321"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "4.321"
 
     mutate_rpc_device_status(
         monkeypatch, mock_rpc_device, "voltmeter:100", "voltage", None
     )
     mock_rpc_device.mock_update()
 
-    assert hass.states.get(entity_id).state == STATE_UNAVAILABLE
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_UNAVAILABLE
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-voltmeter:100-voltmeter"
 
 
@@ -528,15 +542,16 @@ async def test_rpc_polling_sensor(
     entity_id = register_entity(hass, SENSOR_DOMAIN, "test_name_rssi", "wifi-rssi")
     await init_integration(hass, 2)
 
-    assert hass.states.get(entity_id).state == "-63"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "-63"
 
     mutate_rpc_device_status(monkeypatch, mock_rpc_device, "wifi", "rssi", "-70")
     await mock_polling_rpc_update(hass, freezer)
 
-    assert hass.states.get(entity_id).state == "-70"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "-70"
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-wifi-rssi"
 
 
@@ -567,12 +582,14 @@ async def test_rpc_sleeping_sensor(
     mock_rpc_device.mock_online()
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    assert hass.states.get(entity_id).state == "22.9"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "22.9"
 
     mutate_rpc_device_status(monkeypatch, mock_rpc_device, "temperature:0", "tC", 23.4)
     mock_rpc_device.mock_update()
 
-    assert hass.states.get(entity_id).state == "23.4"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "23.4"
 
 
 async def test_rpc_restored_sleeping_sensor(
@@ -600,7 +617,8 @@ async def test_rpc_restored_sleeping_sensor(
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
-    assert hass.states.get(entity_id).state == "21.0"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "21.0"
 
     # Make device online
     monkeypatch.setattr(mock_rpc_device, "initialized", True)
@@ -611,7 +629,8 @@ async def test_rpc_restored_sleeping_sensor(
     mock_rpc_device.mock_update()
     await hass.async_block_till_done()
 
-    assert hass.states.get(entity_id).state == "22.9"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "22.9"
 
 
 async def test_rpc_restored_sleeping_sensor_no_last_state(
@@ -637,7 +656,8 @@ async def test_rpc_restored_sleeping_sensor_no_last_state(
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
-    assert hass.states.get(entity_id).state == STATE_UNKNOWN
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_UNKNOWN
 
     # Make device online
     monkeypatch.setattr(mock_rpc_device, "initialized", True)
@@ -648,7 +668,8 @@ async def test_rpc_restored_sleeping_sensor_no_last_state(
     mock_rpc_device.mock_update()
     await hass.async_block_till_done()
 
-    assert hass.states.get(entity_id).state == "22.9"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "22.9"
 
 
 @pytest.mark.usefixtures("entity_registry_enabled_by_default")
@@ -658,36 +679,32 @@ async def test_rpc_em1_sensors(
     """Test RPC sensors for EM1 component."""
     await init_integration(hass, 2)
 
-    state = hass.states.get("sensor.test_name_em0_power")
+    assert (state := hass.states.get("sensor.test_name_em0_power"))
-    assert state
     assert state.state == "85.3"
 
-    entry = entity_registry.async_get("sensor.test_name_em0_power")
+    assert (entry := entity_registry.async_get("sensor.test_name_em0_power"))
-    assert entry
     assert entry.unique_id == "123456789ABC-em1:0-power_em1"
 
-    state = hass.states.get("sensor.test_name_em1_power")
+    assert (state := hass.states.get("sensor.test_name_em1_power"))
-    assert state
     assert state.state == "123.3"
 
-    entry = entity_registry.async_get("sensor.test_name_em1_power")
+    assert (entry := entity_registry.async_get("sensor.test_name_em1_power"))
-    assert entry
     assert entry.unique_id == "123456789ABC-em1:1-power_em1"
 
-    state = hass.states.get("sensor.test_name_em0_total_active_energy")
+    assert (state := hass.states.get("sensor.test_name_em0_total_active_energy"))
-    assert state
     assert state.state == "123.4564"
 
-    entry = entity_registry.async_get("sensor.test_name_em0_total_active_energy")
+    assert (
-    assert entry
+        entry := entity_registry.async_get("sensor.test_name_em0_total_active_energy")
+    )
     assert entry.unique_id == "123456789ABC-em1data:0-total_act_energy"
 
-    state = hass.states.get("sensor.test_name_em1_total_active_energy")
+    assert (state := hass.states.get("sensor.test_name_em1_total_active_energy"))
-    assert state
     assert state.state == "987.6543"
 
-    entry = entity_registry.async_get("sensor.test_name_em1_total_active_energy")
+    assert (
-    assert entry
+        entry := entity_registry.async_get("sensor.test_name_em1_total_active_energy")
+    )
     assert entry.unique_id == "123456789ABC-em1data:1-total_act_energy"
 
 
@@ -713,7 +730,7 @@ async def test_rpc_sleeping_update_entity_service(
     mock_rpc_device.mock_online()
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == "22.9"
 
     await hass.services.async_call(
@@ -724,11 +741,10 @@ async def test_rpc_sleeping_update_entity_service(
     )
 
     # Entity should be available after update_entity service call
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == "22.9"
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-temperature:0-temperature_0"
 
     assert (
@@ -762,7 +778,8 @@ async def test_block_sleeping_update_entity_service(
     mock_block_device.mock_online()
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    assert hass.states.get(entity_id).state == "22.1"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "22.1"
 
     await hass.services.async_call(
         HA_DOMAIN,
@@ -772,11 +789,10 @@ async def test_block_sleeping_update_entity_service(
     )
 
     # Entity should be available after update_entity service call
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == "22.1"
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-sensor_0-temp"
 
     assert (
@@ -809,20 +825,18 @@ async def test_rpc_analog_input_sensors(
     await init_integration(hass, 2)
 
     entity_id = f"{SENSOR_DOMAIN}.test_name_input_1_analog"
-    assert hass.states.get(entity_id).state == "89"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "89"
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-input:1-analoginput"
 
     entity_id = f"{SENSOR_DOMAIN}.test_name_input_1_analog_value"
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert state
     assert state.state == "8.9"
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == expected_unit
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-input:1-analoginput_xpercent"
 
 
@@ -857,7 +871,8 @@ async def test_rpc_disabled_xpercent(
     await init_integration(hass, 2)
 
     entity_id = f"{SENSOR_DOMAIN}.test_name_input_1_analog"
-    assert hass.states.get(entity_id).state == "89"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "89"
 
     entity_id = f"{SENSOR_DOMAIN}.test_name_input_1_analog_value"
     assert hass.states.get(entity_id) is None
@@ -887,23 +902,20 @@ async def test_rpc_pulse_counter_sensors(
     await init_integration(hass, 2)
 
     entity_id = f"{SENSOR_DOMAIN}.gas_pulse_counter"
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == "56174"
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == "pulse"
     assert state.attributes.get(ATTR_STATE_CLASS) == SensorStateClass.TOTAL
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-input:2-pulse_counter"
 
     entity_id = f"{SENSOR_DOMAIN}.gas_counter_value"
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert state
     assert state.state == "561.74"
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == expected_unit
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-input:2-counter_value"
 
 
@@ -938,7 +950,8 @@ async def test_rpc_disabled_xtotal_counter(
     await init_integration(hass, 2)
 
     entity_id = f"{SENSOR_DOMAIN}.gas_pulse_counter"
-    assert hass.states.get(entity_id).state == "20635"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "20635"
 
     entity_id = f"{SENSOR_DOMAIN}.gas_counter_value"
     assert hass.states.get(entity_id) is None
@@ -968,23 +981,20 @@ async def test_rpc_pulse_counter_frequency_sensors(
     await init_integration(hass, 2)
 
     entity_id = f"{SENSOR_DOMAIN}.gas_pulse_counter_frequency"
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == "208.0"
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == UnitOfFrequency.HERTZ
     assert state.attributes.get(ATTR_STATE_CLASS) == SensorStateClass.MEASUREMENT
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-input:2-counter_frequency"
 
     entity_id = f"{SENSOR_DOMAIN}.gas_pulse_counter_frequency_value"
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert state
     assert state.state == "6.11"
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == expected_unit
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-input:2-counter_frequency_value"
 
 
@@ -1007,11 +1017,9 @@ async def test_rpc_disabled_xfreq(
 
     entity_id = f"{SENSOR_DOMAIN}.gas_pulse_counter_frequency_value"
 
-    state = hass.states.get(entity_id)
+    assert hass.states.get(entity_id) is None
-    assert not state
 
-    entry = entity_registry.async_get(entity_id)
+    assert entity_registry.async_get(entity_id) is None
-    assert not entry
 
 
 @pytest.mark.parametrize(
@@ -1043,17 +1051,16 @@ async def test_rpc_device_virtual_text_sensor(
 
     await init_integration(hass, 3)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert state
     assert state.state == "lorem ipsum"
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-text:203-text"
 
     monkeypatch.setitem(mock_rpc_device.status["text:203"], "value", "dolor sit amet")
     mock_rpc_device.mock_update()
-    assert hass.states.get(entity_id).state == "dolor sit amet"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "dolor sit amet"
 
 
 async def test_rpc_remove_text_virtual_sensor_when_mode_field(
@@ -1086,8 +1093,7 @@ async def test_rpc_remove_text_virtual_sensor_when_mode_field(
     await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
 
-    entry = entity_registry.async_get(entity_id)
+    assert entity_registry.async_get(entity_id) is None
-    assert not entry
 
 
 async def test_rpc_remove_text_virtual_sensor_when_orphaned(
@@ -1111,8 +1117,7 @@ async def test_rpc_remove_text_virtual_sensor_when_orphaned(
     await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
 
-    entry = entity_registry.async_get(entity_id)
+    assert entity_registry.async_get(entity_id) is None
-    assert not entry
 
 
 @pytest.mark.parametrize(
@@ -1148,18 +1153,17 @@ async def test_rpc_device_virtual_number_sensor(
 
     await init_integration(hass, 3)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert state
     assert state.state == "34.5"
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == expected_unit
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-number:203-number"
 
     monkeypatch.setitem(mock_rpc_device.status["number:203"], "value", 56.7)
     mock_rpc_device.mock_update()
-    assert hass.states.get(entity_id).state == "56.7"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "56.7"
 
 
 async def test_rpc_remove_number_virtual_sensor_when_mode_field(
@@ -1197,8 +1201,7 @@ async def test_rpc_remove_number_virtual_sensor_when_mode_field(
     await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
 
-    entry = entity_registry.async_get(entity_id)
+    assert entity_registry.async_get(entity_id) is None
-    assert not entry
 
 
 async def test_rpc_remove_number_virtual_sensor_when_orphaned(
@@ -1222,8 +1225,7 @@ async def test_rpc_remove_number_virtual_sensor_when_orphaned(
     await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
 
-    entry = entity_registry.async_get(entity_id)
+    assert entity_registry.async_get(entity_id) is None
-    assert not entry
 
 
 @pytest.mark.parametrize(
@@ -1263,19 +1265,18 @@ async def test_rpc_device_virtual_enum_sensor(
 
     await init_integration(hass, 3)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert state
     assert state.state == expected_state
     assert state.attributes.get(ATTR_DEVICE_CLASS) == SensorDeviceClass.ENUM
     assert state.attributes.get(ATTR_OPTIONS) == ["Title 1", "two", "three"]
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-enum:203-enum"
 
     monkeypatch.setitem(mock_rpc_device.status["enum:203"], "value", "two")
     mock_rpc_device.mock_update()
-    assert hass.states.get(entity_id).state == "two"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "two"
 
 
 async def test_rpc_remove_enum_virtual_sensor_when_mode_dropdown(
@@ -1317,8 +1318,7 @@ async def test_rpc_remove_enum_virtual_sensor_when_mode_dropdown(
     await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
 
-    entry = entity_registry.async_get(entity_id)
+    assert entity_registry.async_get(entity_id) is None
-    assert not entry
 
 
 async def test_rpc_remove_enum_virtual_sensor_when_orphaned(
@@ -1342,8 +1342,7 @@ async def test_rpc_remove_enum_virtual_sensor_when_orphaned(
     await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
 
-    entry = entity_registry.async_get(entity_id)
+    assert entity_registry.async_get(entity_id) is None
-    assert not entry
 
 
 @pytest.mark.usefixtures("entity_registry_enabled_by_default")
@@ -1374,61 +1373,51 @@ async def test_rpc_rgbw_sensors(
 
     entity_id = f"sensor.test_name_{light_type}_light_0_power"
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert state
     assert state.state == "12.2"
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == UnitOfPower.WATT
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == f"123456789ABC-{light_type}:0-power_{light_type}"
 
     entity_id = f"sensor.test_name_{light_type}_light_0_energy"
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert state
     assert state.state == "0.045141"
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == UnitOfEnergy.KILO_WATT_HOUR
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == f"123456789ABC-{light_type}:0-energy_{light_type}"
 
     entity_id = f"sensor.test_name_{light_type}_light_0_current"
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert state
     assert state.state == "0.23"
     assert (
         state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == UnitOfElectricCurrent.AMPERE
     )
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == f"123456789ABC-{light_type}:0-current_{light_type}"
 
     entity_id = f"sensor.test_name_{light_type}_light_0_voltage"
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert state
     assert state.state == "12.4"
     assert (
         state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == UnitOfElectricPotential.VOLT
     )
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == f"123456789ABC-{light_type}:0-voltage_{light_type}"
 
     entity_id = f"sensor.test_name_{light_type}_light_0_device_temperature"
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert state
     assert state.state == "54.3"
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == UnitOfTemperature.CELSIUS
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == f"123456789ABC-{light_type}:0-temperature_{light_type}"
 
 
@@ -1441,15 +1430,17 @@ async def test_rpc_device_sensor_goes_unavailable_on_disconnect(
 ) -> None:
     """Test RPC device with sensor goes unavailable on disconnect."""
     await init_integration(hass, 2)
-    temp_sensor_state = hass.states.get("sensor.test_name_temperature")
+
-    assert temp_sensor_state is not None
+    assert (state := hass.states.get("sensor.test_name_temperature"))
-    assert temp_sensor_state.state != STATE_UNAVAILABLE
+    assert state.state != STATE_UNAVAILABLE
+
     monkeypatch.setattr(mock_rpc_device, "connected", False)
     monkeypatch.setattr(mock_rpc_device, "initialized", False)
     mock_rpc_device.mock_disconnected()
     await hass.async_block_till_done()
-    temp_sensor_state = hass.states.get("sensor.test_name_temperature")
+
-    assert temp_sensor_state.state == STATE_UNAVAILABLE
+    assert (state := hass.states.get("sensor.test_name_temperature"))
+    assert state.state == STATE_UNAVAILABLE
 
     freezer.tick(60)
     async_fire_time_changed(hass)
@@ -1460,8 +1451,9 @@ async def test_rpc_device_sensor_goes_unavailable_on_disconnect(
     monkeypatch.setattr(mock_rpc_device, "initialized", True)
     mock_rpc_device.mock_initialized()
     await hass.async_block_till_done()
-    temp_sensor_state = hass.states.get("sensor.test_name_temperature")
+
-    assert temp_sensor_state.state != STATE_UNAVAILABLE
+    assert (state := hass.states.get("sensor.test_name_temperature"))
+    assert state.state != STATE_UNAVAILABLE
 
 
 async def test_rpc_voltmeter_value(
@@ -1474,13 +1466,11 @@ async def test_rpc_voltmeter_value(
 
     await init_integration(hass, 2)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-
     assert state.state == "12.34"
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == "ppm"
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-voltmeter:100-voltmeter_value"
 
 
@@ -1525,8 +1515,7 @@ async def test_rpc_device_virtual_number_sensor_with_device_class(
 
     await init_integration(hass, 3)
 
-    state = hass.states.get("sensor.test_name_current_humidity")
+    assert (state := hass.states.get("sensor.test_name_current_humidity"))
-    assert state
     assert state.state == "34"
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == PERCENTAGE
     assert state.attributes.get(ATTR_DEVICE_CLASS) == SensorDeviceClass.HUMIDITY

tests/components/shelly/test_switch.py

@@ -28,7 +28,7 @@ from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.device_registry import DeviceRegistry
 from homeassistant.helpers.entity_registry import EntityRegistry
 
-from . import get_entity_state, init_integration, register_device, register_entity
+from . import init_integration, register_device, register_entity
 
 from tests.common import mock_restore_cache
 
@@ -42,22 +42,25 @@ async def test_block_device_services(
 ) -> None:
     """Test block device turn on/off services."""
     await init_integration(hass, 1)
+    entity_id = "switch.test_name_channel_1"
 
     await hass.services.async_call(
         SWITCH_DOMAIN,
         SERVICE_TURN_ON,
-        {ATTR_ENTITY_ID: "switch.test_name_channel_1"},
+        {ATTR_ENTITY_ID: entity_id},
         blocking=True,
     )
-    assert hass.states.get("switch.test_name_channel_1").state == STATE_ON
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_ON
 
     await hass.services.async_call(
         SWITCH_DOMAIN,
         SERVICE_TURN_OFF,
-        {ATTR_ENTITY_ID: "switch.test_name_channel_1"},
+        {ATTR_ENTITY_ID: entity_id},
         blocking=True,
     )
-    assert hass.states.get("switch.test_name_channel_1").state == STATE_OFF
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_OFF
 
 
 @pytest.mark.parametrize("model", MOTION_MODELS)
@@ -75,7 +78,8 @@ async def test_block_motion_switch(
     mock_block_device.mock_online()
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    assert get_entity_state(hass, entity_id) == STATE_ON
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_ON
 
     # turn off
     await hass.services.async_call(
@@ -88,7 +92,9 @@ async def test_block_motion_switch(
     mock_block_device.mock_update()
 
     mock_block_device.set_shelly_motion_detection.assert_called_once_with(False)
-    assert get_entity_state(hass, entity_id) == STATE_OFF
+
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_OFF
 
     # turn on
     mock_block_device.set_shelly_motion_detection.reset_mock()
@@ -102,7 +108,9 @@ async def test_block_motion_switch(
     mock_block_device.mock_update()
 
     mock_block_device.set_shelly_motion_detection.assert_called_once_with(True)
-    assert get_entity_state(hass, entity_id) == STATE_ON
+
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_ON
 
 
 @pytest.mark.parametrize("model", MOTION_MODELS)
@@ -132,14 +140,16 @@ async def test_block_restored_motion_switch(
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
-    assert get_entity_state(hass, entity_id) == STATE_OFF
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_OFF
 
     # Make device online
     monkeypatch.setattr(mock_block_device, "initialized", True)
     mock_block_device.mock_online()
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    assert get_entity_state(hass, entity_id) == STATE_ON
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_ON
 
 
 @pytest.mark.parametrize("model", MOTION_MODELS)
@@ -167,14 +177,16 @@ async def test_block_restored_motion_switch_no_last_state(
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
-    assert get_entity_state(hass, entity_id) == STATE_UNKNOWN
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_UNKNOWN
 
     # Make device online
     monkeypatch.setattr(mock_block_device, "initialized", True)
     mock_block_device.mock_online()
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    assert get_entity_state(hass, entity_id) == STATE_ON
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_ON
 
 
 @pytest.mark.parametrize(
@@ -205,8 +217,7 @@ async def test_block_device_unique_ids(
         mock_block_device.mock_online()
         await hass.async_block_till_done(wait_background_tasks=True)
 
-    entry = entity_registry.async_get(entity)
+    assert (entry := entity_registry.async_get(entity))
-    assert entry
     assert entry.unique_id == unique_id
 
 
@@ -273,11 +284,15 @@ async def test_block_device_update(
     """Test block device update."""
     monkeypatch.setattr(mock_block_device.blocks[RELAY_BLOCK_ID], "output", False)
     await init_integration(hass, 1)
-    assert hass.states.get("switch.test_name_channel_1").state == STATE_OFF
+
+    entity_id = "switch.test_name_channel_1"
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_OFF
 
     monkeypatch.setattr(mock_block_device.blocks[RELAY_BLOCK_ID], "output", True)
     mock_block_device.mock_update()
-    assert hass.states.get("switch.test_name_channel_1").state == STATE_ON
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_ON
 
 
 async def test_block_device_no_relay_blocks(
@@ -317,23 +332,26 @@ async def test_rpc_device_services(
     monkeypatch.setitem(mock_rpc_device.status["sys"], "relay_in_thermostat", False)
     await init_integration(hass, 2)
 
+    entity_id = "switch.test_switch_0"
     await hass.services.async_call(
         SWITCH_DOMAIN,
         SERVICE_TURN_ON,
-        {ATTR_ENTITY_ID: "switch.test_switch_0"},
+        {ATTR_ENTITY_ID: entity_id},
         blocking=True,
     )
-    assert hass.states.get("switch.test_switch_0").state == STATE_ON
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_ON
 
     monkeypatch.setitem(mock_rpc_device.status["switch:0"], "output", False)
     await hass.services.async_call(
         SWITCH_DOMAIN,
         SERVICE_TURN_OFF,
-        {ATTR_ENTITY_ID: "switch.test_switch_0"},
+        {ATTR_ENTITY_ID: entity_id},
         blocking=True,
     )
     mock_rpc_device.mock_update()
-    assert hass.states.get("switch.test_switch_0").state == STATE_OFF
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_OFF
 
 
 async def test_rpc_device_unique_ids(
@@ -347,8 +365,7 @@ async def test_rpc_device_unique_ids(
     monkeypatch.setitem(mock_rpc_device.status["sys"], "relay_in_thermostat", False)
     await init_integration(hass, 2)
 
-    entry = entity_registry.async_get("switch.test_switch_0")
+    assert (entry := entity_registry.async_get("switch.test_switch_0"))
-    assert entry
     assert entry.unique_id == "123456789ABC-switch:0"
 
 
@@ -360,6 +377,7 @@ async def test_rpc_device_switch_type_lights_mode(
         mock_rpc_device.config["sys"]["ui_data"], "consumption_types", ["lights"]
     )
     await init_integration(hass, 2)
+
     assert hass.states.get("switch.test_switch_0") is None
 
 
@@ -463,7 +481,7 @@ async def test_wall_display_relay_mode(
 
     config_entry = await init_integration(hass, 2, model=MODEL_WALL_DISPLAY)
 
-    assert hass.states.get(climate_entity_id) is not None
+    assert (state := hass.states.get(climate_entity_id))
     assert len(hass.states.async_entity_ids(CLIMATE_DOMAIN)) == 1
 
     new_status = deepcopy(mock_rpc_device.status)
@@ -476,17 +494,16 @@ async def test_wall_display_relay_mode(
     await hass.async_block_till_done()
 
     # the climate entity should be removed
+
     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 := 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 := entity_registry.async_get(switch_entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-switch:0"
 
 
@@ -519,12 +536,10 @@ async def test_rpc_device_virtual_switch(
 
     await init_integration(hass, 3)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert state
     assert state.state == STATE_ON
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-boolean:200-boolean"
 
     monkeypatch.setitem(mock_rpc_device.status["boolean:200"], "value", False)
@@ -535,7 +550,8 @@ async def test_rpc_device_virtual_switch(
         blocking=True,
     )
     mock_rpc_device.mock_update()
-    assert hass.states.get(entity_id).state == STATE_OFF
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_OFF
 
     monkeypatch.setitem(mock_rpc_device.status["boolean:200"], "value", True)
     await hass.services.async_call(
@@ -545,7 +561,8 @@ async def test_rpc_device_virtual_switch(
         blocking=True,
     )
     mock_rpc_device.mock_update()
-    assert hass.states.get(entity_id).state == STATE_ON
+    assert (state := hass.states.get(entity_id))
+    assert state.state == STATE_ON
 
 
 async def test_rpc_device_virtual_binary_sensor(
@@ -566,8 +583,7 @@ async def test_rpc_device_virtual_binary_sensor(
 
     await init_integration(hass, 3)
 
-    state = hass.states.get(entity_id)
+    assert hass.states.get(entity_id) is None
-    assert not state
 
 
 async def test_rpc_remove_virtual_switch_when_mode_label(
@@ -600,8 +616,7 @@ async def test_rpc_remove_virtual_switch_when_mode_label(
     await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
 
-    entry = entity_registry.async_get(entity_id)
+    assert entity_registry.async_get(entity_id) is None
-    assert not entry
 
 
 async def test_rpc_remove_virtual_switch_when_orphaned(
@@ -625,8 +640,7 @@ async def test_rpc_remove_virtual_switch_when_orphaned(
     await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
 
-    entry = entity_registry.async_get(entity_id)
+    assert entity_registry.async_get(entity_id) is None
-    assert not entry
 
 
 @pytest.mark.usefixtures("entity_registry_enabled_by_default")
@@ -656,11 +670,10 @@ async def test_rpc_device_script_switch(
 
     await init_integration(hass, 3)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert state
     assert state.state == STATE_ON
-    entry = entity_registry.async_get(entity_id)
+
-    assert entry
+    assert (entry := entity_registry.async_get(entity_id))
     assert entry.unique_id == f"123456789ABC-{key}-script"
 
     monkeypatch.setitem(mock_rpc_device.status[key], "running", False)
@@ -671,8 +684,8 @@ async def test_rpc_device_script_switch(
         blocking=True,
     )
     mock_rpc_device.mock_update()
-    state = hass.states.get(entity_id)
+
-    assert state
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_OFF
 
     monkeypatch.setitem(mock_rpc_device.status[key], "running", True)
@@ -683,6 +696,6 @@ async def test_rpc_device_script_switch(
         blocking=True,
     )
     mock_rpc_device.mock_update()
-    state = hass.states.get(entity_id)
+
-    assert state
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_ON

tests/components/shelly/test_text.py

@@ -47,17 +47,17 @@ async def test_rpc_device_virtual_text(
 
     await init_integration(hass, 3)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert state
     assert state.state == "lorem ipsum"
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-text:203-text"
 
     monkeypatch.setitem(mock_rpc_device.status["text:203"], "value", "dolor sit amet")
     mock_rpc_device.mock_update()
-    assert hass.states.get(entity_id).state == "dolor sit amet"
+
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "dolor sit amet"
 
     monkeypatch.setitem(mock_rpc_device.status["text:203"], "value", "sed do eiusmod")
     await hass.services.async_call(
@@ -67,7 +67,9 @@ async def test_rpc_device_virtual_text(
         blocking=True,
     )
     mock_rpc_device.mock_update()
-    assert hass.states.get(entity_id).state == "sed do eiusmod"
+
+    assert (state := hass.states.get(entity_id))
+    assert state.state == "sed do eiusmod"
 
 
 async def test_rpc_remove_virtual_text_when_mode_label(
@@ -100,8 +102,7 @@ async def test_rpc_remove_virtual_text_when_mode_label(
     await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
 
-    entry = entity_registry.async_get(entity_id)
+    assert entity_registry.async_get(entity_id) is None
-    assert not entry
 
 
 async def test_rpc_remove_virtual_text_when_orphaned(
@@ -125,5 +126,4 @@ async def test_rpc_remove_virtual_text_when_orphaned(
     await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
 
-    entry = entity_registry.async_get(entity_id)
+    assert entity_registry.async_get(entity_id) is None
-    assert not entry

tests/components/shelly/test_update.py

@@ -61,14 +61,16 @@ async def test_block_update(
     monkeypatch.setitem(mock_block_device.status, "cloud", {"connected": False})
     await init_integration(hass, 1)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_ON
     assert state.attributes[ATTR_INSTALLED_VERSION] == "1.0.0"
     assert state.attributes[ATTR_LATEST_VERSION] == "2.0.0"
     assert state.attributes[ATTR_IN_PROGRESS] is False
     assert state.attributes[ATTR_UPDATE_PERCENTAGE] is None
-    supported_feat = state.attributes[ATTR_SUPPORTED_FEATURES]
+    assert (
-    assert supported_feat == UpdateEntityFeature.INSTALL | UpdateEntityFeature.PROGRESS
+        state.attributes[ATTR_SUPPORTED_FEATURES]
+        == UpdateEntityFeature.INSTALL | UpdateEntityFeature.PROGRESS
+    )
 
     await hass.services.async_call(
         UPDATE_DOMAIN,
@@ -78,7 +80,7 @@ async def test_block_update(
     )
     assert mock_block_device.trigger_ota_update.call_count == 1
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_ON
     assert state.attributes[ATTR_INSTALLED_VERSION] == "1.0.0"
     assert state.attributes[ATTR_LATEST_VERSION] == "2.0.0"
@@ -89,15 +91,14 @@ async def test_block_update(
     monkeypatch.setitem(mock_block_device.status["update"], "old_version", "2.0.0")
     await mock_rest_update(hass, freezer)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_OFF
     assert state.attributes[ATTR_INSTALLED_VERSION] == "2.0.0"
     assert state.attributes[ATTR_LATEST_VERSION] == "2.0.0"
     assert state.attributes[ATTR_IN_PROGRESS] is False
     assert state.attributes[ATTR_UPDATE_PERCENTAGE] is None
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-fwupdate"
 
 
@@ -117,7 +118,7 @@ async def test_block_beta_update(
     monkeypatch.setitem(mock_block_device.status, "cloud", {"connected": False})
     await init_integration(hass, 1)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_OFF
     assert state.attributes[ATTR_INSTALLED_VERSION] == "1.0.0"
     assert state.attributes[ATTR_LATEST_VERSION] == "1.0.0"
@@ -129,7 +130,7 @@ async def test_block_beta_update(
     )
     await mock_rest_update(hass, freezer)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_ON
     assert state.attributes[ATTR_INSTALLED_VERSION] == "1.0.0"
     assert state.attributes[ATTR_LATEST_VERSION] == "2.0.0-beta"
@@ -145,7 +146,7 @@ async def test_block_beta_update(
     )
     assert mock_block_device.trigger_ota_update.call_count == 1
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_ON
     assert state.attributes[ATTR_INSTALLED_VERSION] == "1.0.0"
     assert state.attributes[ATTR_LATEST_VERSION] == "2.0.0-beta"
@@ -155,15 +156,14 @@ async def test_block_beta_update(
     monkeypatch.setitem(mock_block_device.status["update"], "old_version", "2.0.0-beta")
     await mock_rest_update(hass, freezer)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_OFF
     assert state.attributes[ATTR_INSTALLED_VERSION] == "2.0.0-beta"
     assert state.attributes[ATTR_LATEST_VERSION] == "2.0.0-beta"
     assert state.attributes[ATTR_IN_PROGRESS] is False
     assert state.attributes[ATTR_UPDATE_PERCENTAGE] is None
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-fwupdate_beta"
 
 
@@ -256,11 +256,12 @@ async def test_block_version_compare(
     monkeypatch.setitem(mock_block_device.status, "cloud", {"connected": False})
     await init_integration(hass, 1)
 
-    state = hass.states.get(entity_id_latest)
+    assert (state := hass.states.get(entity_id_latest))
     assert state.state == STATE_OFF
     assert state.attributes[ATTR_INSTALLED_VERSION] == STABLE
     assert state.attributes[ATTR_LATEST_VERSION] == STABLE
-    state = hass.states.get(entity_id_beta)
+
+    assert (state := hass.states.get(entity_id_beta))
     assert state.state == STATE_ON
     assert state.attributes[ATTR_INSTALLED_VERSION] == STABLE
     assert state.attributes[ATTR_LATEST_VERSION] == BETA
@@ -270,11 +271,12 @@ async def test_block_version_compare(
     monkeypatch.setitem(mock_block_device.status["update"], "beta_version", BETA)
     await mock_rest_update(hass, freezer)
 
-    state = hass.states.get(entity_id_latest)
+    assert (state := hass.states.get(entity_id_latest))
     assert state.state == STATE_OFF
     assert state.attributes[ATTR_INSTALLED_VERSION] == BETA
     assert state.attributes[ATTR_LATEST_VERSION] == STABLE
-    state = hass.states.get(entity_id_beta)
+
+    assert (state := hass.states.get(entity_id_beta))
     assert state.state == STATE_OFF
     assert state.attributes[ATTR_INSTALLED_VERSION] == BETA
     assert state.attributes[ATTR_LATEST_VERSION] == BETA
@@ -298,7 +300,7 @@ async def test_rpc_update(
     )
     await init_integration(hass, 2)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_ON
     assert state.attributes[ATTR_INSTALLED_VERSION] == "1"
     assert state.attributes[ATTR_LATEST_VERSION] == "2"
@@ -316,7 +318,7 @@ async def test_rpc_update(
 
     assert mock_rpc_device.trigger_ota_update.call_count == 1
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_ON
     assert state.attributes[ATTR_INSTALLED_VERSION] == "1"
     assert state.attributes[ATTR_LATEST_VERSION] == "2"
@@ -339,7 +341,7 @@ async def test_rpc_update(
         },
     )
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.attributes[ATTR_IN_PROGRESS] is True
     assert state.attributes[ATTR_UPDATE_PERCENTAGE] == 0
 
@@ -359,7 +361,7 @@ async def test_rpc_update(
         },
     )
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.attributes[ATTR_IN_PROGRESS] is True
     assert state.attributes[ATTR_UPDATE_PERCENTAGE] == 50
 
@@ -380,15 +382,14 @@ async def test_rpc_update(
     monkeypatch.setitem(mock_rpc_device.shelly, "ver", "2")
     mock_rpc_device.mock_update()
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_OFF
     assert state.attributes[ATTR_INSTALLED_VERSION] == "2"
     assert state.attributes[ATTR_LATEST_VERSION] == "2"
     assert state.attributes[ATTR_IN_PROGRESS] is False
     assert state.attributes[ATTR_UPDATE_PERCENTAGE] is None
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-sys-fwupdate"
 
 
@@ -419,7 +420,7 @@ async def test_rpc_sleeping_update(
     mock_rpc_device.mock_online()
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_ON
     assert state.attributes[ATTR_INSTALLED_VERSION] == "1"
     assert state.attributes[ATTR_LATEST_VERSION] == "2"
@@ -431,7 +432,7 @@ async def test_rpc_sleeping_update(
     monkeypatch.setitem(mock_rpc_device.shelly, "ver", "2")
     mock_rpc_device.mock_update()
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_OFF
     assert state.attributes[ATTR_INSTALLED_VERSION] == "2"
     assert state.attributes[ATTR_LATEST_VERSION] == "2"
@@ -439,8 +440,7 @@ async def test_rpc_sleeping_update(
     assert state.attributes[ATTR_UPDATE_PERCENTAGE] is None
     assert state.attributes[ATTR_SUPPORTED_FEATURES] == UpdateEntityFeature(0)
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-sys-fwupdate"
 
 
@@ -471,7 +471,7 @@ async def test_rpc_restored_sleeping_update(
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_ON
     assert state.attributes[ATTR_INSTALLED_VERSION] == "1"
     assert state.attributes[ATTR_LATEST_VERSION] == "2"
@@ -488,7 +488,7 @@ async def test_rpc_restored_sleeping_update(
     mock_rpc_device.mock_update()
     await hass.async_block_till_done()
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_OFF
     assert state.attributes[ATTR_INSTALLED_VERSION] == "2"
     assert state.attributes[ATTR_LATEST_VERSION] == "2"
@@ -527,7 +527,7 @@ async def test_rpc_restored_sleeping_update_no_last_state(
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_UNKNOWN
 
     # Make device online
@@ -539,7 +539,7 @@ async def test_rpc_restored_sleeping_update_no_last_state(
     mock_rpc_device.mock_update()
     await hass.async_block_till_done()
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_ON
     assert state.attributes[ATTR_INSTALLED_VERSION] == "1"
     assert state.attributes[ATTR_LATEST_VERSION] == "2"
@@ -569,7 +569,7 @@ async def test_rpc_beta_update(
     )
     await init_integration(hass, 2)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_OFF
     assert state.attributes[ATTR_INSTALLED_VERSION] == "1"
     assert state.attributes[ATTR_LATEST_VERSION] == "1"
@@ -586,7 +586,7 @@ async def test_rpc_beta_update(
     )
     await mock_rest_update(hass, freezer)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_ON
     assert state.attributes[ATTR_INSTALLED_VERSION] == "1"
     assert state.attributes[ATTR_LATEST_VERSION] == "2b"
@@ -616,7 +616,7 @@ async def test_rpc_beta_update(
 
     assert mock_rpc_device.trigger_ota_update.call_count == 1
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_ON
     assert state.attributes[ATTR_INSTALLED_VERSION] == "1"
     assert state.attributes[ATTR_LATEST_VERSION] == "2b"
@@ -639,7 +639,7 @@ async def test_rpc_beta_update(
         },
     )
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.attributes[ATTR_IN_PROGRESS] is True
     assert state.attributes[ATTR_UPDATE_PERCENTAGE] == 40
 
@@ -660,15 +660,14 @@ async def test_rpc_beta_update(
     monkeypatch.setitem(mock_rpc_device.shelly, "ver", "2b")
     await mock_rest_update(hass, freezer)
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
     assert state.state == STATE_OFF
     assert state.attributes[ATTR_INSTALLED_VERSION] == "2b"
     assert state.attributes[ATTR_LATEST_VERSION] == "2b"
     assert state.attributes[ATTR_IN_PROGRESS] is False
     assert state.attributes[ATTR_UPDATE_PERCENTAGE] is None
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-sys-fwupdate_beta"
 
 

tests/components/shelly/test_valve.py

@@ -25,11 +25,11 @@ async def test_block_device_gas_valve(
     await init_integration(hass, 1, MODEL_GAS)
     entity_id = "valve.test_name_valve"
 
-    entry = entity_registry.async_get(entity_id)
+    assert (entry := entity_registry.async_get(entity_id))
-    assert entry
     assert entry.unique_id == "123456789ABC-valve_0-valve"
 
-    assert hass.states.get(entity_id).state == ValveState.CLOSED
+    assert (state := hass.states.get(entity_id))
+    assert state.state == ValveState.CLOSED
 
     await hass.services.async_call(
         VALVE_DOMAIN,
@@ -38,16 +38,14 @@ async def test_block_device_gas_valve(
         blocking=True,
     )
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert state
     assert state.state == ValveState.OPENING
 
     monkeypatch.setattr(mock_block_device.blocks[GAS_VALVE_BLOCK_ID], "valve", "opened")
     mock_block_device.mock_update()
     await hass.async_block_till_done()
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert state
     assert state.state == ValveState.OPEN
 
     await hass.services.async_call(
@@ -57,14 +55,12 @@ async def test_block_device_gas_valve(
         blocking=True,
     )
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert state
     assert state.state == ValveState.CLOSING
 
     monkeypatch.setattr(mock_block_device.blocks[GAS_VALVE_BLOCK_ID], "valve", "closed")
     mock_block_device.mock_update()
     await hass.async_block_till_done()
 
-    state = hass.states.get(entity_id)
+    assert (state := hass.states.get(entity_id))
-    assert state
     assert state.state == ValveState.CLOSED