core/tests/components/amberelectric/conftest.py

"""Provide common Amber fixtures."""

from collections.abc import AsyncGenerator, Generator
from unittest.mock import AsyncMock, Mock, patch

from amberelectric.models.interval import Interval
import pytest

from homeassistant.components.amberelectric.const import (
    CONF_SITE_ID,
    CONF_SITE_NAME,
    DOMAIN,
)
from homeassistant.const import CONF_API_TOKEN

from .helpers import (
    CONTROLLED_LOAD_CHANNEL,
    FEED_IN_CHANNEL,
    FORECASTS,
    GENERAL_AND_CONTROLLED_SITE_ID,
    GENERAL_AND_FEED_IN_SITE_ID,
    GENERAL_CHANNEL,
    GENERAL_CHANNEL_WITH_RANGE,
    GENERAL_FORECASTS,
    GENERAL_ONLY_SITE_ID,
)

from tests.common import MockConfigEntry

MOCK_API_TOKEN = "psk_0000000000000000"


def create_amber_config_entry(
    site_id: str, entry_id: str, name: str
) -> MockConfigEntry:
    """Create an Amber config entry."""
    return MockConfigEntry(
        domain=DOMAIN,
        data={
            CONF_API_TOKEN: MOCK_API_TOKEN,
            CONF_SITE_NAME: name,
            CONF_SITE_ID: site_id,
        },
        entry_id=entry_id,
    )


@pytest.fixture
def mock_amber_client() -> Generator[AsyncMock]:
    """Mock the Amber API client."""
    with patch(
        "homeassistant.components.amberelectric.amberelectric.AmberApi",
        autospec=True,
    ) as mock_client:
        yield mock_client


@pytest.fixture
def mock_setup_entry() -> Generator[AsyncMock]:
    """Override async_setup_entry."""
    with patch(
        "homeassistant.components.amberelectric.async_setup_entry", return_value=True
    ) as mock_setup_entry:
        yield mock_setup_entry


@pytest.fixture
async def general_channel_config_entry():
    """Generate the default Amber config entry."""
    return create_amber_config_entry(GENERAL_ONLY_SITE_ID, GENERAL_ONLY_SITE_ID, "home")


@pytest.fixture
async def general_channel_and_controlled_load_config_entry():
    """Generate the default Amber config entry for site with controlled load."""
    return create_amber_config_entry(
        GENERAL_AND_CONTROLLED_SITE_ID, GENERAL_AND_CONTROLLED_SITE_ID, "home"
    )


@pytest.fixture
async def general_channel_and_feed_in_config_entry():
    """Generate the default Amber config entry for site with feed in."""
    return create_amber_config_entry(
        GENERAL_AND_FEED_IN_SITE_ID, GENERAL_AND_FEED_IN_SITE_ID, "home"
    )


@pytest.fixture
def general_channel_prices() -> list[Interval]:
    """List containing general channel prices."""
    return GENERAL_CHANNEL


@pytest.fixture
def general_channel_prices_with_range() -> list[Interval]:
    """List containing general channel prices."""
    return GENERAL_CHANNEL_WITH_RANGE


@pytest.fixture
def controlled_load_channel_prices() -> list[Interval]:
    """List containing controlled load channel prices."""
    return CONTROLLED_LOAD_CHANNEL


@pytest.fixture
def feed_in_channel_prices() -> list[Interval]:
    """List containing feed in channel prices."""
    return FEED_IN_CHANNEL


@pytest.fixture
def forecast_prices() -> list[Interval]:
    """List containing forecasts with advanced prices."""
    return FORECASTS


@pytest.fixture
def general_forecast_prices() -> list[Interval]:
    """List containing forecasts with advanced prices."""
    return GENERAL_FORECASTS


@pytest.fixture
def mock_amber_client_general_channel(
    mock_amber_client: AsyncMock, general_channel_prices: list[Interval]
) -> Generator[AsyncMock]:
    """Fake general channel prices."""
    client = mock_amber_client.return_value
    client.get_current_prices.return_value = general_channel_prices
    return mock_amber_client


@pytest.fixture
def mock_amber_client_general_channel_with_range(
    mock_amber_client: AsyncMock, general_channel_prices_with_range: list[Interval]
) -> Generator[AsyncMock]:
    """Fake general channel prices with a range."""
    client = mock_amber_client.return_value
    client.get_current_prices.return_value = general_channel_prices_with_range
    return mock_amber_client


@pytest.fixture
def mock_amber_client_general_and_controlled_load(
    mock_amber_client: AsyncMock,
    general_channel_prices: list[Interval],
    controlled_load_channel_prices: list[Interval],
) -> Generator[AsyncMock]:
    """Fake general channel and controlled load channel prices."""
    client = mock_amber_client.return_value
    client.get_current_prices.return_value = (
        general_channel_prices + controlled_load_channel_prices
    )
    return mock_amber_client


@pytest.fixture
async def mock_amber_client_general_and_feed_in(
    mock_amber_client: AsyncMock,
    general_channel_prices: list[Interval],
    feed_in_channel_prices: list[Interval],
) -> AsyncGenerator[Mock]:
    """Set up general channel and feed in channel."""
    client = mock_amber_client.return_value
    client.get_current_prices.return_value = (
        general_channel_prices + feed_in_channel_prices
    )
    return mock_amber_client


@pytest.fixture
async def mock_amber_client_forecasts(
    mock_amber_client: AsyncMock, forecast_prices: list[Interval]
) -> AsyncGenerator[Mock]:
    """Set up general channel, controlled load and feed in channel."""
    client = mock_amber_client.return_value
    client.get_current_prices.return_value = forecast_prices
    return mock_amber_client


@pytest.fixture
async def mock_amber_client_general_forecasts(
    mock_amber_client: AsyncMock, general_forecast_prices: list[Interval]
) -> AsyncGenerator[Mock]:
    """Set up general channel only."""
    client = mock_amber_client.return_value
    client.get_current_prices.return_value = general_forecast_prices
    return mock_amber_client