migration test

tests/components/recorder/test_migrate.py

@@ -25,10 +25,15 @@ from homeassistant.components.recorder import db_schema, migration
 from homeassistant.components.recorder.db_schema import (
     SCHEMA_VERSION,
     Events,
+    EventTypes,
     RecorderRuns,
     States,
 )
-from homeassistant.components.recorder.tasks import ContextIDMigrationTask
+from homeassistant.components.recorder.queries import select_event_type_ids
+from homeassistant.components.recorder.tasks import (
+    ContextIDMigrationTask,
+    EventTypeIDMigrationTask,
+)
 from homeassistant.components.recorder.util import session_scope
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import recorder as recorder_helper
@@ -688,3 +693,104 @@ async def test_migrate_context_ids(
     assert invalid_context_id_event["context_id_bin"] == b"\x00" * 16
     assert invalid_context_id_event["context_user_id_bin"] is None
     assert invalid_context_id_event["context_parent_id_bin"] is None
+
+
+@pytest.mark.parametrize("enable_migrate_event_type_ids", [True])
+async def test_migrate_event_type_ids(
+    async_setup_recorder_instance: RecorderInstanceGenerator, hass: HomeAssistant
+) -> None:
+    """Test we can migrate event_types to the EventTypes table."""
+    instance = await async_setup_recorder_instance(hass)
+    await async_wait_recording_done(hass)
+
+    test_uuid = uuid.uuid4()
+    uuid_hex = test_uuid.hex
+
+    def _insert_events():
+        with session_scope(hass=hass) as session:
+            session.add_all(
+                (
+                    Events(
+                        event_type="event_type_one",
+                        event_data=None,
+                        origin_idx=0,
+                        time_fired=None,
+                        time_fired_ts=1677721632.452529,
+                        context_id=uuid_hex,
+                        context_id_bin=None,
+                        context_user_id=None,
+                        context_user_id_bin=None,
+                        context_parent_id=None,
+                        context_parent_id_bin=None,
+                    ),
+                    Events(
+                        event_type="event_type_one",
+                        event_data=None,
+                        origin_idx=0,
+                        time_fired=None,
+                        time_fired_ts=1677721632.552529,
+                        context_id=None,
+                        context_id_bin=None,
+                        context_user_id=None,
+                        context_user_id_bin=None,
+                        context_parent_id=None,
+                        context_parent_id_bin=None,
+                    ),
+                    Events(
+                        event_type="event_type_two",
+                        event_data=None,
+                        origin_idx=0,
+                        time_fired=None,
+                        time_fired_ts=1677721632.552529,
+                        context_id=None,
+                        context_id_bin=None,
+                        context_user_id=None,
+                        context_user_id_bin=None,
+                        context_parent_id=None,
+                        context_parent_id_bin=None,
+                    ),
+                )
+            )
+
+    await instance.async_add_executor_job(_insert_events)
+
+    await async_wait_recording_done(hass)
+    # This is a threadsafe way to add a task to the recorder
+    instance.queue_task(EventTypeIDMigrationTask())
+    await async_recorder_block_till_done(hass)
+
+    def _object_as_dict(obj):
+        return {c.key: getattr(obj, c.key) for c in inspect(obj).mapper.column_attrs}
+
+    def _fetch_migrated_events():
+        with session_scope(hass=hass) as session:
+            events = (
+                session.query(Events.event_id, Events.time_fired, EventTypes.event_type)
+                .filter(
+                    Events.event_type_id.in_(
+                        select_event_type_ids(
+                            (
+                                "event_type_one",
+                                "event_type_two",
+                            )
+                        )
+                    )
+                )
+                .outerjoin(EventTypes, Events.event_type_id == EventTypes.event_type_id)
+                .all()
+            )
+            assert len(events) == 3
+            result = {}
+            for event in events:
+                result.setdefault(event.event_type, []).append(
+                    {
+                        "event_id": event.event_id,
+                        "time_fired": event.time_fired,
+                        "event_type": event.event_type,
+                    }
+                )
+            return result
+
+    events_by_type = await instance.async_add_executor_job(_fetch_migrated_events)
+    assert len(events_by_type["event_type_one"]) == 2
+    assert len(events_by_type["event_type_two"]) == 1

tests/conftest.py

@@ -1148,6 +1148,16 @@ def enable_migrate_context_ids() -> bool:
     return False
 
 
+@pytest.fixture
+def enable_migrate_event_type_ids() -> bool:
+    """Fixture to control enabling of recorder's event type id migration.
+
+    To enable context id migration, tests can be marked with:
+    @pytest.mark.parametrize("enable_migrate_event_type_ids", [True])
+    """
+    return False
+
+
 @pytest.fixture
 def recorder_config() -> dict[str, Any] | None:
     """Fixture to override recorder config.
@@ -1291,6 +1301,7 @@ async def async_setup_recorder_instance(
     enable_statistics: bool,
     enable_statistics_table_validation: bool,
     enable_migrate_context_ids: bool,
+    enable_migrate_event_type_ids: bool,
 ) -> AsyncGenerator[RecorderInstanceGenerator, None]:
     """Yield callable to setup recorder instance."""
     # pylint: disable-next=import-outside-toplevel
@@ -1309,6 +1320,11 @@ async def async_setup_recorder_instance(
     migrate_context_ids = (
         recorder.Recorder._migrate_context_ids if enable_migrate_context_ids else None
     )
+    migrate_event_type_ids = (
+        recorder.Recorder._migrate_event_type_ids
+        if enable_migrate_event_type_ids
+        else None
+    )
     with patch(
         "homeassistant.components.recorder.Recorder.async_nightly_tasks",
         side_effect=nightly,
@@ -1325,6 +1341,10 @@ async def async_setup_recorder_instance(
         "homeassistant.components.recorder.Recorder._migrate_context_ids",
         side_effect=migrate_context_ids,
         autospec=True,
+    ), patch(
+        "homeassistant.components.recorder.Recorder._migrate_event_type_ids",
+        side_effect=migrate_event_type_ids,
+        autospec=True,
     ):
 
         async def async_setup_recorder(