Use attribute enums in integration tests (#181824)

tests/components/accuweather/test_sensor.py

@@ -18,8 +18,8 @@ from homeassistant.components.homeassistant import (
 )
 from homeassistant.const import (
     ATTR_ENTITY_ID,
-    ATTR_UNIT_OF_MEASUREMENT,
     STATE_UNAVAILABLE,
+    EntityStateAttribute,
     Platform,
     UnitOfLength,
     UnitOfSpeed,
@@ -168,18 +168,25 @@ async def test_sensor_imperial_units(
     state = hass.states.get("sensor.home_cloud_ceiling")
     assert state
     assert state.state == "10498.687664042"
-    assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == UnitOfLength.FEET
+    assert (
+        state.attributes.get(EntityStateAttribute.UNIT_OF_MEASUREMENT)
+        == UnitOfLength.FEET
+    )
 
     state = hass.states.get("sensor.home_wind_speed")
     assert state
     assert float(state.state) == pytest.approx(9.00988)
-    assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == UnitOfSpeed.MILES_PER_HOUR
+    assert (
+        state.attributes.get(EntityStateAttribute.UNIT_OF_MEASUREMENT)
+        == UnitOfSpeed.MILES_PER_HOUR
+    )
 
     state = hass.states.get("sensor.home_realfeel_temperature")
     assert state
     assert state.state == "77.18"
     assert (
-        state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == UnitOfTemperature.FAHRENHEIT
+        state.attributes.get(EntityStateAttribute.UNIT_OF_MEASUREMENT)
+        == UnitOfTemperature.FAHRENHEIT
     )
 
 

tests/components/airpatrol/test_climate.py

@@ -27,6 +27,7 @@ from homeassistant.components.climate import (
     SERVICE_SET_TEMPERATURE,
     SWING_OFF,
     SWING_ON,
+    ClimateEntityStateAttribute,
     HVACMode,
 )
 from homeassistant.const import (
@@ -122,7 +123,7 @@ async def test_climate_set_temperature(
     TARGET_TEMP = 25.0
 
     state = hass.states.get("climate.living_room")
-    assert state.attributes[ATTR_TEMPERATURE] == 22.0
+    assert state.attributes[ClimateEntityStateAttribute.TARGET_TEMPERATURE] == 22.0
 
     climate_data["ParametersData"]["PumpTemp"] = f"{TARGET_TEMP:.3f}"
     await hass.services.async_call(
@@ -136,7 +137,9 @@ async def test_climate_set_temperature(
 
     get_client.set_unit_climate_data.assert_called_once()
     state = hass.states.get("climate.living_room")
-    assert state.attributes[ATTR_TEMPERATURE] == TARGET_TEMP
+    assert (
+        state.attributes[ClimateEntityStateAttribute.TARGET_TEMPERATURE] == TARGET_TEMP
+    )
 
 
 async def test_climate_set_hvac_mode(
@@ -172,7 +175,7 @@ async def test_climate_set_fan_mode(
 ) -> None:
     """Test setting fan mode."""
     state = hass.states.get("climate.living_room")
-    assert state.attributes[ATTR_FAN_MODE] == FAN_HIGH
+    assert state.attributes[ClimateEntityStateAttribute.FAN_MODE] == FAN_HIGH
 
     climate_data["ParametersData"]["FanSpeed"] = HA_TO_AP_FAN_MODES[FAN_LOW]
     await hass.services.async_call(
@@ -186,7 +189,7 @@ async def test_climate_set_fan_mode(
 
     get_client.set_unit_climate_data.assert_called_once()
     state = hass.states.get("climate.living_room")
-    assert state.attributes[ATTR_FAN_MODE] == FAN_LOW
+    assert state.attributes[ClimateEntityStateAttribute.FAN_MODE] == FAN_LOW
 
 
 async def test_climate_set_swing_mode(
@@ -197,7 +200,7 @@ async def test_climate_set_swing_mode(
 ) -> None:
     """Test setting swing mode."""
     state = hass.states.get("climate.living_room")
-    assert state.attributes[ATTR_SWING_MODE] == SWING_OFF
+    assert state.attributes[ClimateEntityStateAttribute.SWING_MODE] == SWING_OFF
 
     climate_data["ParametersData"]["Swing"] = HA_TO_AP_SWING_MODES[SWING_ON]
     await hass.services.async_call(
@@ -211,7 +214,7 @@ async def test_climate_set_swing_mode(
 
     get_client.set_unit_climate_data.assert_called_once()
     state = hass.states.get("climate.living_room")
-    assert state.attributes[ATTR_SWING_MODE] == SWING_ON
+    assert state.attributes[ClimateEntityStateAttribute.SWING_MODE] == SWING_ON
 
 
 @pytest.mark.parametrize(
@@ -313,7 +316,7 @@ async def test_climate_set_temperature_api_error(
 ) -> None:
     """Test async_set_temperature handles API error."""
     state = hass.states.get("climate.living_room")
-    assert state.attributes[ATTR_TEMPERATURE] == 22.0
+    assert state.attributes[ClimateEntityStateAttribute.TARGET_TEMPERATURE] == 22.0
 
     get_client.set_unit_climate_data.side_effect = Exception("API Error")
 
@@ -327,7 +330,7 @@ async def test_climate_set_temperature_api_error(
     )
 
     state = hass.states.get("climate.living_room")
-    assert state.attributes[ATTR_TEMPERATURE] == 22.0
+    assert state.attributes[ClimateEntityStateAttribute.TARGET_TEMPERATURE] == 22.0
 
 
 @pytest.mark.parametrize(
@@ -354,7 +357,7 @@ async def test_climate_fan_mode_invalid(
 ) -> None:
     """Test fan_mode with unexpected value."""
     state = hass.states.get("climate.living_room")
-    assert state.attributes[ATTR_FAN_MODE] is None
+    assert state.attributes[ClimateEntityStateAttribute.FAN_MODE] is None
 
 
 @pytest.mark.parametrize(
@@ -380,4 +383,4 @@ async def test_climate_swing_mode_invalid(
 ) -> None:
     """Test swing_mode with unexpected value."""
     state = hass.states.get("climate.living_room")
-    assert state.attributes[ATTR_SWING_MODE] is None
+    assert state.attributes[ClimateEntityStateAttribute.SWING_MODE] is None

tests/components/airtouch5/test_cover.py

@@ -10,12 +10,12 @@ from airtouch5py.packets.zone_status import (
 from syrupy.assertion import SnapshotAssertion
 
 from homeassistant.components.cover import (
-    ATTR_CURRENT_POSITION,
     ATTR_POSITION,
     DOMAIN as COVER_DOMAIN,
     SERVICE_CLOSE_COVER,
     SERVICE_OPEN_COVER,
     SERVICE_SET_COVER_POSITION,
+    CoverEntityStateAttribute,
     CoverState,
 )
 from homeassistant.const import ATTR_ENTITY_ID, Platform
@@ -95,7 +95,9 @@ async def test_cover_callbacks(
     zone_2_initial = hass.states.get(COVER_ZONE_2_ENTITY_ID)
     assert zone_2_initial
     zone_2_initial_state = zone_2_initial.state
-    zone_2_initial_position = zone_2_initial.attributes.get(ATTR_CURRENT_POSITION)
+    zone_2_initial_position = zone_2_initial.attributes.get(
+        CoverEntityStateAttribute.CURRENT_POSITION
+    )
 
     # Define a method to call all zone_status_callbacks, as the real client would
     async def _call_zone_status_callback(open_percentage: float) -> None:
@@ -122,37 +124,49 @@ async def test_cover_callbacks(
     state = hass.states.get(COVER_ENTITY_ID)
     assert state
     assert state.state == CoverState.OPEN
-    assert state.attributes.get(ATTR_CURRENT_POSITION) == 70
+    assert state.attributes.get(CoverEntityStateAttribute.CURRENT_POSITION) == 70
     zone_2 = hass.states.get(COVER_ZONE_2_ENTITY_ID)
     assert zone_2 and zone_2.state == zone_2_initial_state
-    assert zone_2.attributes.get(ATTR_CURRENT_POSITION) == zone_2_initial_position
+    assert (
+        zone_2.attributes.get(CoverEntityStateAttribute.CURRENT_POSITION)
+        == zone_2_initial_position
+    )
 
     # Fully open
     await _call_zone_status_callback(1)
     state = hass.states.get(COVER_ENTITY_ID)
     assert state
     assert state.state == CoverState.OPEN
-    assert state.attributes.get(ATTR_CURRENT_POSITION) == 100
+    assert state.attributes.get(CoverEntityStateAttribute.CURRENT_POSITION) == 100
     zone_2 = hass.states.get(COVER_ZONE_2_ENTITY_ID)
     assert zone_2 and zone_2.state == zone_2_initial_state
-    assert zone_2.attributes.get(ATTR_CURRENT_POSITION) == zone_2_initial_position
+    assert (
+        zone_2.attributes.get(CoverEntityStateAttribute.CURRENT_POSITION)
+        == zone_2_initial_position
+    )
 
     # Fully closed
     await _call_zone_status_callback(0.0)
     state = hass.states.get(COVER_ENTITY_ID)
     assert state
     assert state.state == CoverState.CLOSED
-    assert state.attributes.get(ATTR_CURRENT_POSITION) == 0
+    assert state.attributes.get(CoverEntityStateAttribute.CURRENT_POSITION) == 0
     zone_2 = hass.states.get(COVER_ZONE_2_ENTITY_ID)
     assert zone_2 and zone_2.state == zone_2_initial_state
-    assert zone_2.attributes.get(ATTR_CURRENT_POSITION) == zone_2_initial_position
+    assert (
+        zone_2.attributes.get(CoverEntityStateAttribute.CURRENT_POSITION)
+        == zone_2_initial_position
+    )
 
     # Partly reopened
     await _call_zone_status_callback(0.3)
     state = hass.states.get(COVER_ENTITY_ID)
     assert state
     assert state.state == CoverState.OPEN
-    assert state.attributes.get(ATTR_CURRENT_POSITION) == 30
+    assert state.attributes.get(CoverEntityStateAttribute.CURRENT_POSITION) == 30
     zone_2 = hass.states.get(COVER_ZONE_2_ENTITY_ID)
     assert zone_2 and zone_2.state == zone_2_initial_state
-    assert zone_2.attributes.get(ATTR_CURRENT_POSITION) == zone_2_initial_position
+    assert (
+        zone_2.attributes.get(CoverEntityStateAttribute.CURRENT_POSITION)
+        == zone_2_initial_position
+    )

tests/components/aosmith/test_water_heater.py

@@ -20,7 +20,7 @@ from homeassistant.components.water_heater import (
     STATE_HEAT_PUMP,
     WaterHeaterEntityFeature,
 )
-from homeassistant.const import ATTR_ENTITY_ID, ATTR_SUPPORTED_FEATURES, Platform
+from homeassistant.const import ATTR_ENTITY_ID, EntityStateAttribute, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import entity_registry as er
@@ -59,7 +59,7 @@ async def test_state_away_mode_unsupported(
     """Test away mode unsupported if water heater lacks vacation mode."""
     state = hass.states.get("water_heater.basement_my_water_heater")
     assert (
-        state.attributes.get(ATTR_SUPPORTED_FEATURES)
+        state.attributes.get(EntityStateAttribute.SUPPORTED_FEATURES)
         == WaterHeaterEntityFeature.TARGET_TEMPERATURE
         | WaterHeaterEntityFeature.OPERATION_MODE
     )

tests/components/arcam_fmj/test_media_player.py

@@ -16,14 +16,11 @@ from homeassistant.components.homeassistant import (
 )
 from homeassistant.components.media_player import (
     ATTR_INPUT_SOURCE,
-    ATTR_MEDIA_ARTIST,
-    ATTR_MEDIA_CHANNEL,
     ATTR_MEDIA_CONTENT_ID,
     ATTR_MEDIA_CONTENT_TYPE,
     ATTR_MEDIA_VOLUME_LEVEL,
     ATTR_MEDIA_VOLUME_MUTED,
     ATTR_SOUND_MODE,
-    ATTR_SOUND_MODE_LIST,
     DOMAIN as MEDIA_PLAYER_DOMAIN,
     SERVICE_PLAY_MEDIA,
     SERVICE_SELECT_SOUND_MODE,
@@ -34,6 +31,8 @@ from homeassistant.components.media_player import (
     SERVICE_VOLUME_MUTE,
     SERVICE_VOLUME_SET,
     SERVICE_VOLUME_UP,
+    MediaPlayerEntityCapabilityAttribute,
+    MediaPlayerEntityStateAttribute,
     MediaType,
 )
 from homeassistant.const import ATTR_ENTITY_ID, STATE_UNAVAILABLE, Platform
@@ -375,7 +374,7 @@ async def test_sound_mode(
     """Test selection sound mode."""
     state_1.get_decode_mode.return_value = mode_enum
     data = await update(hass, client, MOCK_ENTITY_ID)
-    assert data.attributes.get(ATTR_SOUND_MODE) == mode
+    assert data.attributes.get(MediaPlayerEntityStateAttribute.SOUND_MODE) == mode
 
 
 @pytest.mark.parametrize(
@@ -397,7 +396,10 @@ async def test_sound_mode_list(
     """Test sound mode list."""
     state_1.get_decode_modes.return_value = modes_enum
     data = await update(hass, client, MOCK_ENTITY_ID)
-    assert data.attributes.get(ATTR_SOUND_MODE_LIST) == modes
+    assert (
+        data.attributes.get(MediaPlayerEntityCapabilityAttribute.SOUND_MODE_LIST)
+        == modes
+    )
 
 
 @pytest.mark.usefixtures("player_setup")
@@ -407,15 +409,21 @@ async def test_is_volume_muted(
     """Test muted."""
     state_1.get_mute.return_value = True
     data = await update(hass, client, MOCK_ENTITY_ID)
-    assert data.attributes.get(ATTR_MEDIA_VOLUME_MUTED) is True
+    assert (
+        data.attributes.get(MediaPlayerEntityStateAttribute.MEDIA_VOLUME_MUTED) is True
+    )
 
     state_1.get_mute.return_value = False
     data = await update(hass, client, MOCK_ENTITY_ID)
-    assert data.attributes.get(ATTR_MEDIA_VOLUME_MUTED) is False
+    assert (
+        data.attributes.get(MediaPlayerEntityStateAttribute.MEDIA_VOLUME_MUTED) is False
+    )
 
     state_1.get_mute.return_value = None
     data = await update(hass, client, MOCK_ENTITY_ID)
-    assert data.attributes.get(ATTR_MEDIA_VOLUME_MUTED) is None
+    assert (
+        data.attributes.get(MediaPlayerEntityStateAttribute.MEDIA_VOLUME_MUTED) is None
+    )
 
 
 @pytest.mark.usefixtures("player_setup")
@@ -423,15 +431,21 @@ async def test_volume_level(hass: HomeAssistant, client: Mock, state_1: State) -
     """Test volume."""
     state_1.get_volume.return_value = 0
     data = await update(hass, client, MOCK_ENTITY_ID)
-    assert isclose(data.attributes[ATTR_MEDIA_VOLUME_LEVEL], 0.0)
+    assert isclose(
+        data.attributes[MediaPlayerEntityStateAttribute.MEDIA_VOLUME_LEVEL], 0.0
+    )
 
     state_1.get_volume.return_value = 50
     data = await update(hass, client, MOCK_ENTITY_ID)
-    assert isclose(data.attributes[ATTR_MEDIA_VOLUME_LEVEL], 50.0 / 99)
+    assert isclose(
+        data.attributes[MediaPlayerEntityStateAttribute.MEDIA_VOLUME_LEVEL], 50.0 / 99
+    )
 
     state_1.get_volume.return_value = 99
     data = await update(hass, client, MOCK_ENTITY_ID)
-    assert isclose(data.attributes[ATTR_MEDIA_VOLUME_LEVEL], 1.0)
+    assert isclose(
+        data.attributes[MediaPlayerEntityStateAttribute.MEDIA_VOLUME_LEVEL], 1.0
+    )
 
     state_1.get_volume.return_value = None
     data = await update(hass, client, MOCK_ENTITY_ID)
@@ -495,7 +509,10 @@ async def test_media_content_type(
     """Test content type deduction."""
     state_1.get_source.return_value = source
     data = await update(hass, client, MOCK_ENTITY_ID)
-    assert data.attributes.get(ATTR_MEDIA_CONTENT_TYPE) == media_content_type
+    assert (
+        data.attributes.get(MediaPlayerEntityStateAttribute.MEDIA_CONTENT_TYPE)
+        == media_content_type
+    )
 
 
 @pytest.mark.parametrize(
@@ -523,7 +540,7 @@ async def test_media_channel(
     state_1.get_rds_information.return_value = rds
     state_1.get_source.return_value = source
     data = await update(hass, client, MOCK_ENTITY_ID)
-    assert data.attributes.get(ATTR_MEDIA_CHANNEL) == channel
+    assert data.attributes.get(MediaPlayerEntityStateAttribute.MEDIA_CHANNEL) == channel
 
 
 @pytest.mark.parametrize(
@@ -547,7 +564,7 @@ async def test_media_artist(
     state_1.get_dls_pdt.return_value = dls
     state_1.get_source.return_value = source
     data = await update(hass, client, MOCK_ENTITY_ID)
-    assert data.attributes.get(ATTR_MEDIA_ARTIST) == artist
+    assert data.attributes.get(MediaPlayerEntityStateAttribute.MEDIA_ARTIST) == artist
 
 
 @pytest.mark.parametrize(

tests/components/august/test_sensor.py

@@ -3,12 +3,7 @@
 from typing import Any
 
 from homeassistant import core as ha
-from homeassistant.const import (
-    ATTR_ENTITY_PICTURE,
-    ATTR_UNIT_OF_MEASUREMENT,
-    PERCENTAGE,
-    STATE_UNKNOWN,
-)
+from homeassistant.const import PERCENTAGE, STATE_UNKNOWN, EntityStateAttribute
 from homeassistant.core import CoreState, HomeAssistant
 from homeassistant.helpers import entity_registry as er
 
@@ -84,7 +79,10 @@ async def test_create_lock_with_linked_keypad(
         "sensor.front_front_door_lock_keypad_battery"
     )
     assert keypad_battery_state.state == "62"
-    assert keypad_battery_state.attributes[ATTR_UNIT_OF_MEASUREMENT] == PERCENTAGE
+    assert (
+        keypad_battery_state.attributes[EntityStateAttribute.UNIT_OF_MEASUREMENT]
+        == PERCENTAGE
+    )
     entry = entity_registry.async_get("sensor.front_front_door_lock_keypad_battery")
     assert entry
     assert entry.unique_id == "5bc65c24e6ef2a263e1450a8_linked_keypad_battery"
@@ -111,7 +109,10 @@ async def test_create_lock_with_low_battery_linked_keypad(
 
     keypad_battery_state = states.get("sensor.front_front_door_lock_keypad_battery")
     assert keypad_battery_state.state == "10"
-    assert keypad_battery_state.attributes[ATTR_UNIT_OF_MEASUREMENT] == PERCENTAGE
+    assert (
+        keypad_battery_state.attributes[EntityStateAttribute.UNIT_OF_MEASUREMENT]
+        == PERCENTAGE
+    )
     entry = entity_registry.async_get("sensor.front_front_door_lock_keypad_battery")
     assert entry
     assert entry.unique_id == "5bc65c24e6ef2a263e1450a8_linked_keypad_battery"
@@ -339,7 +340,7 @@ async def test_restored_state(
             "keypad": False,
             "tag": True,
             "autorelock": False,
-            ATTR_ENTITY_PICTURE: "image.png",
+            EntityStateAttribute.ENTITY_PICTURE: "image.png",
         },
     )
 
@@ -362,4 +363,4 @@ async def test_restored_state(
     state = hass.states.get(entity_id)
     assert state.state == "Tag Unlock"
     assert state.attributes["method"] == "tag"
-    assert state.attributes[ATTR_ENTITY_PICTURE] == "image.png"
+    assert state.attributes[EntityStateAttribute.ENTITY_PICTURE] == "image.png"

tests/components/axis/test_event.py

@@ -5,9 +5,9 @@ from unittest.mock import patch
 import pytest
 
 from homeassistant.components.event import (
-    ATTR_EVENT_TYPE,
-    ATTR_EVENT_TYPES,
     DOMAIN as EVENT_DOMAIN,
+    EventEntityCapabilityAttribute,
+    EventEntityStateAttribute,
 )
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
@@ -41,7 +41,7 @@ async def test_doorbell_event_entity_created_and_triggered(
 
     state = hass.states.get(event_entities[0])
     assert state is not None
-    assert state.attributes[ATTR_EVENT_TYPES] == ["ring"]
+    assert state.attributes[EventEntityCapabilityAttribute.EVENT_TYPES] == ["ring"]
 
     mock_rtsp_event(
         topic="tns1:Device/tnsaxis:IO/Port",
@@ -55,7 +55,7 @@ async def test_doorbell_event_entity_created_and_triggered(
 
     state = hass.states.get(event_entities[0])
     assert state is not None
-    assert state.attributes[ATTR_EVENT_TYPE] == "ring"
+    assert state.attributes[EventEntityStateAttribute.EVENT_TYPE] == "ring"
 
 
 @pytest.mark.parametrize(