From 47cbda5c5e484a203328184f2d871b1722cd2107 Mon Sep 17 00:00:00 2001
From: Bas Nijholt <basnijholt@gmail.com>
Date: Sat, 6 Apr 2024 19:05:18 +0200
Subject: [PATCH] Add versions to matrix and fix area test (#967)

---
 .github/workflows/pytest.yaml |  4 +++
 test_dependencies.py          |  1 +
 tests/test_switch.py          | 61 ++++++++++++++++++++++++++++++++---
 3 files changed, 61 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/pytest.yaml b/.github/workflows/pytest.yaml
index 750fba04..59a46c1c 100644
--- a/.github/workflows/pytest.yaml
+++ b/.github/workflows/pytest.yaml
@@ -46,6 +46,10 @@ jobs:
             core-version: "2024.1.6"
           - python-version: "3.11"
             core-version: "2024.2.5"
+          - python-version: "3.12"
+            core-version: "2024.3.3"
+          - python-version: "3.12"
+            core-version: "2024.4.1"
           - python-version: "3.12"
             core-version: "dev"
     steps:
diff --git a/test_dependencies.py b/test_dependencies.py
index 96377cee..b92ff1fb 100644
--- a/test_dependencies.py
+++ b/test_dependencies.py
@@ -39,5 +39,6 @@ required = [
     "components.ffmpeg",  # needed since 2024.1
 ]
 to_install = [package for r in required for package in deps[r]]
+to_install.append("flaky")
 
 print(" ".join(to_install))  # noqa: T201
diff --git a/tests/test_switch.py b/tests/test_switch.py
index 9c2f254f..a6c4a2a8 100644
--- a/tests/test_switch.py
+++ b/tests/test_switch.py
@@ -5,6 +5,7 @@ import asyncio
 import contextlib
 import datetime
 import logging
+from collections import OrderedDict
 from copy import deepcopy
 from random import randint
 from typing import Any
@@ -15,6 +16,7 @@ import homeassistant.util.dt as dt_util
 import pytest
 import ulid_transform
 import voluptuous.error
+from flaky import flaky
 from homeassistant.components.adaptive_lighting.adaptation_utils import (
     AdaptationData,
     _create_service_call_data_iterator,
@@ -96,13 +98,15 @@ from homeassistant.const import (
     STATE_OFF,
     STATE_ON,
 )
+from homeassistant.const import __version__ as ha_version
 from homeassistant.core import Context, Event, HomeAssistant, State
+from homeassistant.helpers import area_registry as ar
 from homeassistant.helpers import entity_registry
 from homeassistant.helpers.entity_platform import async_get_platforms
 from homeassistant.setup import async_setup_component
 from homeassistant.util.color import color_temperature_mired_to_kelvin
 
-from tests.common import MockConfigEntry, mock_area_registry
+from tests.common import MockConfigEntry
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -774,6 +778,7 @@ async def test_manual_control(
     assert not manual_control[ENTITY_LIGHT_1]
 
 
+@flaky(max_runs=3, min_passes=1)
 async def test_auto_reset_manual_control(hass):
     switch, (light, *_) = await setup_lights_and_switch(
         hass,
@@ -1328,11 +1333,57 @@ async def test_separate_turn_on_commands(hass, separate_turn_on_commands):
     assert sleep_color_temp != color_temp
 
 
+# Vendored in this function as it was broken
+# https://github.com/home-assistant/core/pull/112150 (my PR and reported issue)
+# Then removed: https://github.com/home-assistant/core/pull/112172
+# Then re-added: https://github.com/home-assistant/core/pull/113453
+# This version is no longer the same as the one in HA because of the many changes
+# that have been made in 2024.
+def mock_area_registry(
+    hass: HomeAssistant,
+) -> ar.AreaRegistry:
+    """Mock the Area Registry."""
+    registry = ar.AreaRegistry(hass)
+    registry._area_data = {}
+    area_kwargs = {
+        "name": "Test Area",
+        "normalized_name": "test-area",
+        "id": "test-area",
+        "picture": None,
+    }
+    year, month = (int(x) for x in ha_version.split(".")[:2])
+    dt = datetime.date(year, month, 1)
+    if dt >= datetime.date(2023, 1, 1):
+        area_kwargs["aliases"] = {}
+    if dt >= datetime.date(2024, 2, 1):
+        area_kwargs["icon"] = None
+    if dt >= datetime.date(2024, 3, 1):
+        area_kwargs["floor_id"] = "test-floor"
+
+    # This mess... 🤯
+    if dt >= datetime.date(2024, 2, 1) and dt != datetime.date(2024, 4, 1):
+        # 2024.4 removed AreaRegistryItems and then added it back in 2024.5:
+        # https://github.com/home-assistant/core/pull/114777
+        registry.areas = ar.AreaRegistryItems()
+    elif dt == datetime.date(2024, 4, 1):
+        from homeassistant.helpers.normalized_name_base_registry import (
+            NormalizedNameBaseRegistryItems,
+        )
+
+        registry.areas = NormalizedNameBaseRegistryItems()
+    else:
+        registry.areas = OrderedDict()
+
+    area = ar.AreaEntry(**area_kwargs)
+    registry.areas[area.id] = area
+    hass.data[ar.DATA_REGISTRY] = registry
+    return registry
+
+
 async def test_light_switch_in_specific_area(hass):
     switch, (light, *_) = await setup_lights_and_switch(hass)
 
-    area_registry = mock_area_registry(hass)
-    area_registry.async_create("test_area")
+    mock_area_registry(hass)
 
     entity = entity_registry.async_get(hass).async_get_or_create(
         LIGHT_DOMAIN,
@@ -1341,9 +1392,9 @@ async def test_light_switch_in_specific_area(hass):
     )
     entity = entity_registry.async_get(hass).async_update_entity(
         entity.entity_id,
-        area_id="test_area",
+        area_id="test-area",
     )
-    _LOGGER.debug("test_area entity: %s", entity)
+    _LOGGER.debug("test-area entity: %s", entity)
     await hass.services.async_call(
         LIGHT_DOMAIN,
         SERVICE_TURN_ON,