Repository: home-assistant/core
File: tests/components/surepetcare/test_sensor.py (Python)

Code before:
DEFAULT_DISABLED_ENTITIES = [
    "sensor.pet_last_seen_flap_device_id",
    "sensor.pet_last_seen_user_id",
]


@pytest.mark.usefixtures("entity_registry_enabled_by_default")
async def test_sensors(
    hass: HomeAssistant,
    entity_registry: er.EntityRegistry,
    surepetcare,
    mock_config_entry_setup: MockConfigEntry,
) -> None:
    """Test the generation of unique ids and sensor states."""
    state_entity_ids = hass.states.async_entity_ids()

    for entity_id, unique_id in EXPECTED_ENTITY_IDS.items():
        assert entity_id in state_entity_ids
        state = hass.states.get(entity_id)
        assert state
        if entity_id == "sensor.pet_last_seen_flap_device_id":
            assert state.state == str(MOCK_PET["position"]["device_id"])
        elif entity_id == "sensor.pet_last_seen_user_id":
            assert state.state == str(MOCK_PET["position"]["user_id"])
        elif entity_id.endswith("_battery_level"):
            assert state.state == "100"
        entity = entity_registry.async_get(entity_id)
        assert entity.unique_id == unique_id


async def test_default_disabled_sensors(
    hass: HomeAssistant,
    entity_registry: er.EntityRegistry,
    surepetcare,
    mock_config_entry_setup: MockConfigEntry,
) -> None:
    """Test sensor entities that are disabled by default."""
    for entity_id in DEFAULT_DISABLED_ENTITIES:
        entity = entity_registry.async_get(entity_id)
        assert entity
        assert entity.disabled_by == RegistryEntryDisabler.INTEGRATION
        assert not hass.states.get(entity_id)


Review comment (refactor):
we don't want if statements in the tests. You could also consider to use the `snapshot_platform` helper here, but that might need a little work.

Other option could be to make the dict a tuple with expected state as well

Code after:
    (
        "sensor.pet_last_seen_flap_device_id",
        f"{HOUSEHOLD_ID}-24680-last_seen_flap_device",
        str(MOCK_PET["position"]["device_id"]),
    ),
    (
        "sensor.pet_last_seen_user_id",
        f"{HOUSEHOLD_ID}-24680-last_seen_user",
        str(MOCK_PET["position"]["user_id"]),
    ),
)

DEFAULT_DISABLED_ENTITIES = [
    "sensor.pet_last_seen_flap_device_id",
    "sensor.pet_last_seen_user_id",
]


@pytest.mark.usefixtures("entity_registry_enabled_by_default")
async def test_sensors(
    hass: HomeAssistant,
    entity_registry: er.EntityRegistry,
    surepetcare,
    mock_config_entry_setup: MockConfigEntry,
) -> None:
    """Test the generation of unique ids and sensor states."""
    state_entity_ids = hass.states.async_entity_ids()

    for entity_id, unique_id, expected_state in EXPECTED_ENTITIES:
        assert entity_id in state_entity_ids
        state = hass.states.get(entity_id)
        assert state
        assert state.state == expected_state
        entity = entity_registry.async_get(entity_id)
        assert entity.unique_id == unique_id


async def test_default_disabled_sensors(
    hass: HomeAssistant,
    entity_registry: er.EntityRegistry,
    surepetcare,
    mock_config_entry_setup: MockConfigEntry,
) -> None:
    """Test sensor entities that are disabled by default."""
    for entity_id in DEFAULT_DISABLED_ENTITIES:
        entity = entity_registry.async_get(entity_id)
        assert entity
        assert entity.disabled_by == RegistryEntryDisabler.INTEGRATION
        assert not hass.states.get(entity_id)