core/tests/components/energy/conftest.py

"""Fixtures for energy component tests."""

from unittest.mock import patch

import pytest

from homeassistant.components.energy import async_get_manager
from homeassistant.components.energy.data import EnergyManager
from homeassistant.components.recorder import Recorder
from homeassistant.core import HomeAssistant
from homeassistant.setup import async_setup_component


@pytest.fixture
def mock_is_entity_recorded():
    """Mock recorder.is_entity_recorded."""
    mocks = {}

    with patch(
        "homeassistant.components.recorder.is_entity_recorded",
        side_effect=lambda hass, entity_id: mocks.get(entity_id, True),
    ):
        yield mocks


@pytest.fixture
def mock_get_metadata():
    """Mock recorder.statistics.get_metadata."""
    mocks = {}

    def _get_metadata(_hass, *, statistic_ids):
        result = {}
        for statistic_id in statistic_ids:
            if statistic_id in mocks:
                if mocks[statistic_id] is not None:
                    result[statistic_id] = mocks[statistic_id]
            else:
                result[statistic_id] = (1, {})
        return result

    with patch(
        "homeassistant.components.recorder.statistics.get_metadata",
        wraps=_get_metadata,
    ):
        yield mocks


@pytest.fixture
async def mock_energy_manager(
    recorder_mock: Recorder, hass: HomeAssistant
) -> EnergyManager:
    """Set up energy."""
    assert await async_setup_component(hass, "energy", {"energy": {}})
    manager = await async_get_manager(hass)
    manager.data = manager.default_preferences()
    return manager