Repository: home-assistant/core
File: tests/components/eheimdigital/test_init.py (Python)

Code before:
        "00:00:00:00:00:02"
    ].heater_data
    eheimdigital_hub_mock.return_value.devices["00:00:00:00:00:02"].heater_data = None
    with (
        patch(
            "homeassistant.components.eheimdigital.coordinator.asyncio.Event",
            new=AsyncMock,
        ),
    ):
        await hass.config_entries.async_setup(mock_config_entry.entry_id)

    for device in eheimdigital_hub_mock.return_value.devices:
        await eheimdigital_hub_mock.call_args.kwargs["device_found_callback"](
            device, eheimdigital_hub_mock.return_value.devices[device].device_type
        )
        await hass.async_block_till_done()

    assert "00:00:00:00:00:02" in mock_config_entry.runtime_data.incomplete_devices

    eheimdigital_hub_mock.return_value.devices[
        "00:00:00:00:00:02"
    ].heater_data = heater_data

    await mock_config_entry.runtime_data._async_receive_callback()

    assert len(mock_config_entry.runtime_data.incomplete_devices) == 0
    assert hass.states.get("number.mock_heater_night_temperature_offset").state == str(
        eheimdigital_hub_mock.return_value.devices[
            "00:00:00:00:00:02"
        ].night_temperature_offset
    )


async def test_remove_device(
    hass: HomeAssistant,
    eheimdigital_hub_mock: MagicMock,
    mock_config_entry: MockConfigEntry,
    device_registry: dr.DeviceRegistry,
    hass_ws_client: WebSocketGenerator,
) -> None:
    """Test removing a device."""
    assert await async_setup_component(hass, "config", {})

    await init_integration(hass, mock_config_entry)

    await eheimdigital_hub_mock.call_args.kwargs["device_found_callback"](
        "00:00:00:00:00:01", EheimDeviceType.VERSION_EHEIM_CLASSIC_LED_CTRL_PLUS_E
    )
    await hass.async_block_till_done()

    mac_address: str = eheimdigital_hub_mock.return_value.main.mac_address

Review comment (suggestion):
Let's avoid touching runtime_data as that is an internal. I think it makes sense to check if the device/entities does not exists after receiving incomplete data and then when we recieve the data we can assert it exists now

Code after:
    mock_config_entry.add_to_hass(hass)
    heater_data = eheimdigital_hub_mock.return_value.devices[
        "00:00:00:00:00:02"
    ].heater_data
    eheimdigital_hub_mock.return_value.devices["00:00:00:00:00:02"].heater_data = None
    with (
        patch(
            "homeassistant.components.eheimdigital.coordinator.asyncio.Event",
            new=AsyncMock,
        ),
    ):
        await hass.config_entries.async_setup(mock_config_entry.entry_id)

    for device in eheimdigital_hub_mock.return_value.devices:
        await eheimdigital_hub_mock.call_args.kwargs["device_found_callback"](
            device, eheimdigital_hub_mock.return_value.devices[device].device_type
        )
        await hass.async_block_till_done()

    assert (
        entity_registry.async_get_entity_id(
            DOMAIN, Platform.NUMBER, "mock_heater_night_temperature_offset"
        )
        is None
    )

    eheimdigital_hub_mock.return_value.devices[
        "00:00:00:00:00:02"
    ].heater_data = heater_data

    await eheimdigital_hub_mock.call_args.kwargs["receive_callback"]()

    assert hass.states.get("number.mock_heater_night_temperature_offset").state == str(
        eheimdigital_hub_mock.return_value.devices[
            "00:00:00:00:00:02"
        ].night_temperature_offset
    )


async def test_remove_device(
    hass: HomeAssistant,
    eheimdigital_hub_mock: MagicMock,
    mock_config_entry: MockConfigEntry,
    device_registry: dr.DeviceRegistry,
    hass_ws_client: WebSocketGenerator,
) -> None:
    """Test removing a device."""
    assert await async_setup_component(hass, "config", {})

    await init_integration(hass, mock_config_entry)