core/tests/components/calendar/conftest.py

"""Test fixtures for calendar sensor platforms."""

import datetime
import secrets
from typing import Any
from unittest.mock import AsyncMock

import pytest
from typing_extensions import Generator

from homeassistant.components.calendar import DOMAIN, CalendarEntity, CalendarEvent
from homeassistant.config_entries import ConfigEntry, ConfigFlow
from homeassistant.const import Platform
from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.util import dt as dt_util

from tests.common import (
    MockConfigEntry,
    MockModule,
    MockPlatform,
    mock_config_flow,
    mock_integration,
    mock_platform,
)

TEST_DOMAIN = "test"


@pytest.fixture
async def set_time_zone(hass: HomeAssistant) -> None:
    """Set the time zone for the tests."""
    # Set our timezone to CST/Regina so we can check calculations
    # This keeps UTC-6 all year round
    await hass.config.async_set_time_zone("America/Regina")


class MockFlow(ConfigFlow):
    """Test flow."""


class MockCalendarEntity(CalendarEntity):
    """Test Calendar entity."""

    _attr_has_entity_name = True

    def __init__(self, name: str, events: list[CalendarEvent] | None = None) -> None:
        """Initialize entity."""
        self._attr_name = name.capitalize()
        self._events = events or []

    @property
    def event(self) -> CalendarEvent | None:
        """Return the next upcoming event."""
        return self._events[0] if self._events else None

    def create_event(
        self,
        start: datetime.datetime,
        end: datetime.datetime,
        summary: str | None = None,
        description: str | None = None,
        location: str | None = None,
    ) -> dict[str, Any]:
        """Create a new fake event, used by tests."""
        event = CalendarEvent(
            start=start,
            end=end,
            summary=summary if summary else f"Event {secrets.token_hex(16)}",
            description=description,
            location=location,
        )
        self._events.append(event)
        return event.as_dict()

    async def async_get_events(
        self,
        hass: HomeAssistant,
        start_date: datetime.datetime,
        end_date: datetime.datetime,
    ) -> list[CalendarEvent]:
        """Return calendar events within a datetime range."""
        assert start_date < end_date
        events = []
        for event in self._events:
            if event.start_datetime_local >= end_date:
                continue
            if event.end_datetime_local < start_date:
                continue
            events.append(event)
        return events


@pytest.fixture
def config_flow_fixture(hass: HomeAssistant) -> Generator[None]:
    """Mock config flow."""
    mock_platform(hass, f"{TEST_DOMAIN}.config_flow")

    with mock_config_flow(TEST_DOMAIN, MockFlow):
        yield


@pytest.fixture(name="config_entry")
async def mock_config_entry(hass: HomeAssistant) -> MockConfigEntry:
    """Create a mock config entry."""
    config_entry = MockConfigEntry(domain=TEST_DOMAIN)
    config_entry.add_to_hass(hass)
    return config_entry


@pytest.fixture
def mock_setup_integration(
    hass: HomeAssistant,
    config_flow_fixture: None,
    test_entities: list[CalendarEntity],
) -> None:
    """Fixture to set up a mock integration."""

    async def async_setup_entry_init(
        hass: HomeAssistant, config_entry: ConfigEntry
    ) -> bool:
        """Set up test config entry."""
        await hass.config_entries.async_forward_entry_setups(config_entry, [DOMAIN])
        return True

    async def async_unload_entry_init(
        hass: HomeAssistant,
        config_entry: ConfigEntry,
    ) -> bool:
        await hass.config_entries.async_unload_platforms(
            config_entry, [Platform.CALENDAR]
        )
        return True

    mock_platform(hass, f"{TEST_DOMAIN}.config_flow")
    mock_integration(
        hass,
        MockModule(
            TEST_DOMAIN,
            async_setup_entry=async_setup_entry_init,
            async_unload_entry=async_unload_entry_init,
        ),
    )

    async def async_setup_entry_platform(
        hass: HomeAssistant,
        config_entry: ConfigEntry,
        async_add_entities: AddEntitiesCallback,
    ) -> None:
        """Set up test event platform via config entry."""
        new_entities = create_test_entities()
        test_entities.clear()
        test_entities.extend(new_entities)
        async_add_entities(test_entities)

    mock_platform(
        hass,
        f"{TEST_DOMAIN}.{DOMAIN}",
        MockPlatform(async_setup_entry=async_setup_entry_platform),
    )


@pytest.fixture(name="test_entities")
def mock_test_entities() -> list[MockCalendarEntity]:
    """Fixture that holdes the fake entities created during the test."""
    return []


def create_test_entities() -> list[MockCalendarEntity]:
    """Create test entities used during the test."""
    half_hour_from_now = dt_util.now() + datetime.timedelta(minutes=30)
    entity1 = MockCalendarEntity(
        "Calendar 1",
        [
            CalendarEvent(
                start=half_hour_from_now,
                end=half_hour_from_now + datetime.timedelta(minutes=60),
                summary="Future Event",
                description="Future Description",
                location="Future Location",
            )
        ],
    )
    entity1.async_get_events = AsyncMock(wraps=entity1.async_get_events)

    middle_of_event = dt_util.now() - datetime.timedelta(minutes=30)
    entity2 = MockCalendarEntity(
        "Calendar 2",
        [
            CalendarEvent(
                start=middle_of_event,
                end=middle_of_event + datetime.timedelta(minutes=60),
                summary="Current Event",
            )
        ],
    )
    entity2.async_get_events = AsyncMock(wraps=entity2.async_get_events)

    return [entity1, entity2]